From ba68efe0f966a5e6fdab28b4056ee5fe0fa72e98 Mon Sep 17 00:00:00 2001 From: "gordon.hamilton" Date: Wed, 15 Jul 2026 13:54:33 -0400 Subject: [PATCH 1/7] feat: implementing window function expressions --- GRAMMAR.md | 65 ++++++ src/parser/expression_grammar.pest | 35 ++- src/parser/expressions.rs | 342 ++++++++++++++++++++++++++++- src/textify/expressions.rs | 334 +++++++++++++++++++++++++++- tests/plan_roundtrip.rs | 86 ++++++++ 5 files changed, 845 insertions(+), 17 deletions(-) diff --git a/GRAMMAR.md b/GRAMMAR.md index 5ba706ef..0fa62465 100644 --- a/GRAMMAR.md +++ b/GRAMMAR.md @@ -501,6 +501,71 @@ Root[status] # assert_eq!(plan.relations.len(), 1); ``` +### Window Functions + +A window function computes a value over a "window" of rows related to the current row (partitioning, ordering, and an optional frame), rather than collapsing rows the way an aggregate does. + +#### Syntax + +`window_function := function_signature anchor? urn_anchor? "(" (expression ("," expression)*)? ")" "over(" window_named_arg ("," window_named_arg)* ")" ":" type` + +`window_named_arg` is one of: + +- `partition=(expression, ...)` - partitioning expressions +- `order=sort_field` or `order=(sort_field, sort_field, ...)` - ordering field(s); a single sort field is written bare (e.g. `order=($1,&AscNullsLast)`), two or more are wrapped in a parenthesized list of tuples (e.g. `order=(($1,&AscNullsLast),($2,&DescNullsLast))`) +- `invocation=&Distinct` / `invocation=&All` - aggregation invocation +- `phase=&InitialToResult` (etc.) - aggregation phase +- `rows=(lower, upper)` / `range=(lower, upper)` - the window frame, each bound is an integer (negative = preceding, + positive = following, `0` = current row) or `_` for unbounded + +`partition=`, `order=`, `invocation=`, and `rows=`/`range=` are all +optional and are omitted when empty; `phase=` is always required and +always printed. A `range=` frame requires exactly one `order=` field. + +#### Examples + +```rust +# use substrait_explain::Parser; +# +# let plan_text = r#" +=== Extensions +URNs: + @ 1: https://github.com/substrait-io/substrait/blob/main/extensions/functions_aggregate.yaml +Functions: + ## 10 @ 1: sum + +=== Plan +Root[a, b, s] + Project[$0, $1, sum($1) over(phase=&InitialToResult, order=($1,&AscNullsLast), invocation=&Distinct, partition=($0), rows=(-3, 0)):fp64?] + Read[t => a:i32, b:fp64] +# "#; +# +# let plan = Parser::parse(plan_text).unwrap(); +# assert_eq!(plan.relations.len(), 1); +``` + +Ranking-style window functions have no meaningful frame, so the `rows=`/`range=` clause is omitted entirely: + +```rust +# use substrait_explain::Parser; +# +# let plan_text = r#" +=== Extensions +URNs: + @ 1: https://github.com/substrait-io/substrait/blob/main/extensions/functions_arithmetic.yaml +Functions: + ## 10 @ 1: row_number + +=== Plan +Root[a, r] + Project[$0, row_number() over(phase=&InitialToResult, order=($0,&AscNullsLast), partition=($0)):i64] + Read[t => a:i32] +# "#; +# +# let plan = Parser::parse(plan_text).unwrap(); +# assert_eq!(plan.relations.len(), 1); +``` + ## Relations Relations represent the operations in a query plan. Each relation is displayed on a single line with indentation showing the hierarchy. diff --git a/src/parser/expression_grammar.pest b/src/parser/expression_grammar.pest index d47408d3..66135b03 100644 --- a/src/parser/expression_grammar.pest +++ b/src/parser/expression_grammar.pest @@ -179,10 +179,43 @@ cast_failure_behavior = { "?" | "!" } // (78:i32)::!i16 THROW_EXCEPTION cast_expression = { "(" ~ sp ~ expression ~ sp ~ ")" ~ sp ~ "::" ~ sp ~ cast_failure_behavior? ~ sp ~ type } +// -- Window Function Calls -- +// +// WindowFunction[Expression.WindowFunction]: a function_call followed +// by a "over(...)" clause carrying window-specific named arguments. +// Example: sum($0) over(partition=($1,$2), order=($3,&AscNullsLast), rows=(-3, 0), phase=&InitialToResult, invocation=&Distinct):fp64 +// +// - `partition=(...)` - optional list of partitioning expressions +// - `order=...` - optional sort key(s); a single sort field is written bare +// (e.g. `order=($3,&AscNullsLast)`), two or more are wrapped in a list +// (e.g. `order=(($3,&AscNullsLast),($4,&DescNullsLast))`) +// - `rows=(lower,upper)` / `range=(lower,upper)` - optional frame bounds, encodes `bounds_type` with rows and range. +// Each bound is an integer offset (negative preceding, positive following, 0 for the current row) +// or `_` for unbounded/unspecified. +// - `phase=&...` - required aggregation phase +// - `invocation=&...` - optional invocation (e.g. `&All`, `&Distinct`) +window_partition_arg = { "partition" ~ sp ~ "=" ~ sp ~ "(" ~ sp ~ expression_list? ~ sp ~ ")" } +order_value = { sort_field | ("(" ~ sp ~ sort_field ~ (sp ~ "," ~ sp ~ sort_field)+ ~ sp ~ ")") } +window_order_arg = { "order" ~ sp ~ "=" ~ sp ~ order_value } +window_bound = { integer | empty } +window_frame_kind = { "rows" | "range" } +window_frame_arg = { window_frame_kind ~ sp ~ "=" ~ sp ~ "(" ~ sp ~ window_bound ~ sp ~ "," ~ sp ~ window_bound ~ sp ~ ")" } +window_phase_arg = { "phase" ~ sp ~ "=" ~ sp ~ enum_value } +window_invocation_arg = { "invocation" ~ sp ~ "=" ~ sp ~ enum_value } + +window_named_arg = { window_partition_arg | window_order_arg | window_frame_arg | window_phase_arg | window_invocation_arg } +window_named_arg_list = { (window_named_arg ~ (sp ~ "," ~ sp ~ window_named_arg)*)? } + +window_function_call = { + function_signature ~ sp ~ anchor? ~ sp ~ urn_anchor? ~ sp ~ argument_list ~ sp ~ "over" ~ sp ~ "(" ~ sp ~ window_named_arg_list ~ sp ~ ")" ~ ":" ~ sp ~ type +} + // Top-level Expression Rule // Order matters for PEGs: Since an identifier can be a function call, we put that first. // cast_expression must come before literal/function_call since it starts with "(" -expression = { if_then | cast_expression | function_call | reference | literal } +// window_function_call must come before function_call: both share the same +// prefix, but window_function_call requires the trailing "over(...)" clause. +expression = { if_then | cast_expression | window_function_call | function_call | reference | literal } // == Extensions == // These rules are for parsing extension declarations, by line. diff --git a/src/parser/expressions.rs b/src/parser/expressions.rs index ff3b153b..fafc48e8 100644 --- a/src/parser/expressions.rs +++ b/src/parser/expressions.rs @@ -1,15 +1,19 @@ use chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime}; +use substrait::proto::aggregate_function::AggregationInvocation; use substrait::proto::aggregate_rel::Measure; use substrait::proto::expression::field_reference::{ReferenceType, RootReference, RootType}; use substrait::proto::expression::if_then::IfClause; use substrait::proto::expression::literal::LiteralType; +use substrait::proto::expression::window_function::{self, bound}; use substrait::proto::expression::{ - Cast, FieldReference, IfThen, Literal, ReferenceSegment, RexType, ScalarFunction, cast, - reference_segment, + Cast, FieldReference, IfThen, Literal, ReferenceSegment, RexType, ScalarFunction, + WindowFunction, cast, reference_segment, }; use substrait::proto::function_argument::ArgType; use substrait::proto::r#type::{Fp64, I64, Kind, Nullability}; -use substrait::proto::{AggregateFunction, Expression, FunctionArgument, Type}; +use substrait::proto::{ + AggregateFunction, AggregationPhase, Expression, FunctionArgument, SortField, Type, +}; use super::types::get_and_validate_anchor; use super::{ @@ -458,6 +462,225 @@ impl ScopedParsePair for ScalarFunction { } } +/// Parse a `window_bound` pair (`integer | empty`) into a `window_function::Bound`. +/// `_` (empty) means unbounded, represented explicitly as `Some(Bound { kind: +/// Some(Unbounded{}) })` rather than `None`. `0` maps to `CurrentRow`; +/// negative/positive integers map to `Preceding`/`Following`. +fn parse_window_bound(pair: pest::iterators::Pair) -> Option { + assert_eq!(pair.as_rule(), Rule::window_bound); + let inner = unwrap_single_pair(pair); + match inner.as_rule() { + Rule::empty => Some(window_function::Bound { + kind: Some(bound::Kind::Unbounded(bound::Unbounded {})), + }), + Rule::integer => { + let offset: i64 = inner.as_str().parse().unwrap(); + let kind = match offset { + 0 => bound::Kind::CurrentRow(bound::CurrentRow {}), + n if n > 0 => bound::Kind::Following(bound::Following { offset: n }), + n => bound::Kind::Preceding(bound::Preceding { offset: -n }), + }; + Some(window_function::Bound { kind: Some(kind) }) + } + other => unreachable!("Grammar guarantees window_bound is integer or empty, got {other:?}"), + } +} + +/// Parse an `enum_value` pair (`&Identifier`) into an `AggregationPhase`. +fn parse_aggregation_phase(pair: pest::iterators::Pair) -> Result { + assert_eq!(pair.as_rule(), Rule::enum_value); + let name = pair.as_str().trim_start_matches('&'); + let phase = match name { + "Unspecified" => AggregationPhase::Unspecified, + "InitialToIntermediate" => AggregationPhase::InitialToIntermediate, + "IntermediateToIntermediate" => AggregationPhase::IntermediateToIntermediate, + "InitialToResult" => AggregationPhase::InitialToResult, + "IntermediateToResult" => AggregationPhase::IntermediateToResult, + other => { + return Err(MessageParseError::invalid( + "AggregationPhase", + pair.as_span(), + format!("Unknown aggregation phase: {other}"), + )); + } + }; + Ok(phase as i32) +} + +/// Parse an `enum_value` pair (`&Identifier`) into an `AggregationInvocation`. +fn parse_aggregation_invocation( + pair: pest::iterators::Pair, +) -> Result { + assert_eq!(pair.as_rule(), Rule::enum_value); + let name = pair.as_str().trim_start_matches('&'); + let invocation = match name { + "Unspecified" => AggregationInvocation::Unspecified, + "All" => AggregationInvocation::All, + "Distinct" => AggregationInvocation::Distinct, + other => { + return Err(MessageParseError::invalid( + "AggregationInvocation", + pair.as_span(), + format!("Unknown aggregation invocation: {other}"), + )); + } + }; + Ok(invocation as i32) +} + +impl ScopedParsePair for WindowFunction { + fn rule() -> Rule { + Rule::window_function_call + } + + fn message() -> &'static str { + "WindowFunction" + } + + fn parse_pair( + extensions: &SimpleExtensions, + pair: pest::iterators::Pair, + ) -> Result { + assert_eq!(pair.as_rule(), Self::rule()); + let span = pair.as_span(); + let mut iter = RuleIter::from(pair.into_inner()); + + // Parse compound function name (required) — e.g. "row_number" or "sum:i64" + let name = iter.parse_next::(); + + // Parse optional anchor (e.g., #1) + let anchor = iter + .try_pop(Rule::anchor) + .map(|n| unwrap_single_pair(n).as_str().parse::().unwrap()); + + // Parse optional URN anchor (e.g., @1) + let _urn_anchor = iter + .try_pop(Rule::urn_anchor) + .map(|n| unwrap_single_pair(n).as_str().parse::().unwrap()); + + // Parse argument list (required) + let argument_list = iter.pop(Rule::argument_list); + let mut arguments = Vec::new(); + for e in argument_list.into_inner() { + arguments.push(FunctionArgument { + arg_type: Some(ArgType::Value(Expression::parse_pair(extensions, e)?)), + }); + } + + // Parse the required `over(...)` named-argument list + let named_arg_list = iter.pop(Rule::window_named_arg_list); + + // Parse required output type (e.g., :i64) + let output_type = Some(Type::parse_pair(extensions, iter.pop(Rule::r#type))?); + iter.done(); + + let mut partitions = Vec::new(); + let mut sorts = Vec::new(); + let mut invocation = AggregationInvocation::Unspecified as i32; + let mut phase = None; + let mut bounds_type = window_function::BoundsType::Unspecified as i32; + let mut lower_bound = None; + let mut upper_bound = None; + + for arg in named_arg_list.into_inner() { + assert_eq!(arg.as_rule(), Rule::window_named_arg); + let inner = unwrap_single_pair(arg); + match inner.as_rule() { + Rule::window_partition_arg => { + let mut parts_iter = RuleIter::from(inner.into_inner()); + if let Some(expr_list) = parts_iter.try_pop(Rule::expression_list) { + for e in expr_list.into_inner() { + partitions.push(Expression::parse_pair(extensions, e)?); + } + } + parts_iter.done(); + } + Rule::window_order_arg => { + let mut order_iter = RuleIter::from(inner.into_inner()); + let order_value = order_iter.pop(Rule::order_value); + order_iter.done(); + for sf in order_value.into_inner() { + sorts.push(SortField::parse_pair(extensions, sf)?); + } + } + Rule::window_frame_arg => { + let mut frame_iter = RuleIter::from(inner.into_inner()); + let kind = frame_iter.pop(Rule::window_frame_kind); + let lower = frame_iter.pop(Rule::window_bound); + let upper = frame_iter.pop(Rule::window_bound); + frame_iter.done(); + bounds_type = match kind.as_str() { + "rows" => window_function::BoundsType::Rows as i32, + "range" => window_function::BoundsType::Range as i32, + other => unreachable!( + "Grammar guarantees window_frame_kind is rows or range, got {other:?}" + ), + }; + lower_bound = parse_window_bound(lower); + upper_bound = parse_window_bound(upper); + } + Rule::window_phase_arg => { + let mut phase_iter = RuleIter::from(inner.into_inner()); + let enum_pair = phase_iter.pop(Rule::enum_value); + phase_iter.done(); + phase = Some(parse_aggregation_phase(enum_pair)?); + } + Rule::window_invocation_arg => { + let mut inv_iter = RuleIter::from(inner.into_inner()); + let enum_pair = inv_iter.pop(Rule::enum_value); + inv_iter.done(); + invocation = parse_aggregation_invocation(enum_pair)?; + } + other => { + unreachable!("Grammar guarantees window_named_arg alternatives, got {other:?}") + } + } + } + + let phase = phase.ok_or_else(|| { + MessageParseError::invalid( + "WindowFunction", + span, + "Missing required phase= argument in over(...)", + ) + })?; + + if bounds_type == window_function::BoundsType::Range as i32 && sorts.len() != 1 { + return Err(MessageParseError::invalid( + "WindowFunction", + span, + format!( + "range= frame requires exactly one order= field, got {}", + sorts.len() + ), + )); + } + + let anchor = get_and_validate_anchor( + extensions, + ExtensionKind::Function, + anchor, + name.full(), + span, + )?; + Ok(WindowFunction { + function_reference: anchor, + arguments, + options: vec![], // TODO: Function Options + output_type, + phase, + sorts, + invocation, + partitions, + bounds_type, + lower_bound, + upper_bound, + #[allow(deprecated)] + args: vec![], + }) + } +} + impl ScopedParsePair for Cast { fn rule() -> Rule { Rule::cast_expression @@ -526,6 +749,11 @@ impl ScopedParsePair for Expression { extensions, inner, )?)), }), + Rule::window_function_call => Ok(Expression { + rex_type: Some(RexType::WindowFunction(WindowFunction::parse_pair( + extensions, inner, + )?)), + }), Rule::reference => Ok(Expression { rex_type: Some(RexType::Selection(Box::new(FieldReference::parse_pair( inner, @@ -542,7 +770,7 @@ impl ScopedParsePair for Expression { )?))), }), _ => unreachable!( - "Grammar guarantees expression can only be literal, function_call, reference, if_then, or cast_expression, got: {:?}", + "Grammar guarantees expression can only be literal, function_call, window_function_call, reference, if_then, or cast_expression, got: {:?}", inner.as_rule() ), } @@ -1450,4 +1678,110 @@ mod tests { "u! prefix in function call base name must be rejected by the grammar" ); } + + #[test] + fn test_window_function_full() { + let exts = make_extensions_for_fn_tests(); + let pair = parse_exact( + Rule::window_function_call, + "add:i64_i64($0, $1) over(partition=($0), order=($1,&AscNullsLast), invocation=&Distinct, rows=(-3, 0), phase=&InitialToResult):i64", + ); + let f = WindowFunction::parse_pair(&exts, pair).unwrap(); + assert_eq!(f.function_reference, 3); + assert_eq!(f.arguments.len(), 2); + assert_eq!(f.partitions.len(), 1); + assert_eq!(f.sorts.len(), 1); + assert_eq!(f.invocation, AggregationInvocation::Distinct as i32); + assert_eq!(f.phase, AggregationPhase::InitialToResult as i32); + assert_eq!(f.bounds_type, window_function::BoundsType::Rows as i32); + assert_eq!( + f.lower_bound, + Some(window_function::Bound { + kind: Some(bound::Kind::Preceding(bound::Preceding { offset: 3 })), + }) + ); + assert_eq!( + f.upper_bound, + Some(window_function::Bound { + kind: Some(bound::Kind::CurrentRow(bound::CurrentRow {})), + }) + ); + } + + #[test] + fn test_window_function_unbounded() { + let exts = make_extensions_for_fn_tests(); + let pair = parse_exact( + Rule::window_function_call, + "add:i64_i64($0, $1) over(phase=&InitialToResult):i64", + ); + let f = WindowFunction::parse_pair(&exts, pair).unwrap(); + assert_eq!( + f.bounds_type, + window_function::BoundsType::Unspecified as i32 + ); + assert_eq!(f.lower_bound, None); + assert_eq!(f.upper_bound, None); + assert_eq!(f.invocation, AggregationInvocation::Unspecified as i32); + assert!(f.partitions.is_empty()); + assert!(f.sorts.is_empty()); + } + + #[test] + fn test_window_function_unbounded_lower_bound() { + let exts = make_extensions_for_fn_tests(); + let pair = parse_exact( + Rule::window_function_call, + "add:i64_i64($0, $1) over(rows=(_, 5), phase=&InitialToResult):i64", + ); + let f = WindowFunction::parse_pair(&exts, pair).unwrap(); + assert_eq!( + f.lower_bound, + Some(window_function::Bound { + kind: Some(bound::Kind::Unbounded(bound::Unbounded {})), + }) + ); + assert_eq!( + f.upper_bound, + Some(window_function::Bound { + kind: Some(bound::Kind::Following(bound::Following { offset: 5 })), + }) + ); + } + + #[test] + fn test_window_function_missing_phase_fails() { + let exts = make_extensions_for_fn_tests(); + let pair = parse_exact( + Rule::window_function_call, + "add:i64_i64($0, $1) over(partition=($0)):i64", + ); + let result = WindowFunction::parse_pair(&exts, pair); + assert!(result.is_err(), "missing phase= must be rejected"); + } + + #[test] + fn test_window_function_range_with_multiple_order_fields_fails() { + let exts = make_extensions_for_fn_tests(); + let pair = parse_exact( + Rule::window_function_call, + "add:i64_i64($0, $1) over(order=(($0,&AscNullsLast),($1,&AscNullsLast)), range=(_, 0), phase=&InitialToResult):i64", + ); + let result = WindowFunction::parse_pair(&exts, pair); + assert!( + result.is_err(), + "range= with more than one order= field must be rejected" + ); + } + + #[test] + fn test_window_function_call_in_expression() { + let exts = make_extensions_for_fn_tests(); + let pair = parse_exact( + Rule::expression, + "add:i64_i64($0, $1) over(phase=&InitialToResult):i64", + ); + let e = Expression::parse_pair(&exts, pair).unwrap(); + assert!(matches!(e.rex_type, Some(RexType::WindowFunction(_)))); + } } diff --git a/src/textify/expressions.rs b/src/textify/expressions.rs index 7d9aaa67..8a738956 100644 --- a/src/textify/expressions.rs +++ b/src/textify/expressions.rs @@ -1,15 +1,20 @@ -use std::fmt::{self}; +use std::fmt::{self, Write as _}; use chrono::{DateTime, NaiveDate}; use expr::RexType; +use substrait::proto::aggregate_function::AggregationInvocation; use substrait::proto::expression::field_reference::{ReferenceType, RootReference, RootType}; use substrait::proto::expression::literal::LiteralType; +use substrait::proto::expression::window_function::{self, bound}; use substrait::proto::expression::{ - Cast, FieldReference, IfThen, ReferenceSegment, ScalarFunction, cast, reference_segment, + Cast, FieldReference, IfThen, ReferenceSegment, ScalarFunction, WindowFunction, cast, + reference_segment, }; use substrait::proto::function_argument::ArgType; +use substrait::proto::sort_field::{SortDirection, SortKind}; use substrait::proto::{ - AggregateFunction, Expression, FunctionArgument, FunctionOption, expression as expr, + AggregateFunction, AggregationPhase, Expression, FunctionArgument, FunctionOption, SortField, + expression as expr, }; use super::{PlanError, Scope, Textify, Visibility}; @@ -522,6 +527,204 @@ impl Textify for IfThen { } } +/// Resolve a `SortField`'s direction to its enum variant name, for `order=`. +fn sort_direction_str(sk: Option<&SortKind>) -> Result<&'static str, PlanError> { + match sk { + Some(SortKind::Direction(d)) => match SortDirection::try_from(*d) { + Ok(SortDirection::AscNullsFirst) => Ok("AscNullsFirst"), + Ok(SortDirection::AscNullsLast) => Ok("AscNullsLast"), + Ok(SortDirection::DescNullsFirst) => Ok("DescNullsFirst"), + Ok(SortDirection::DescNullsLast) => Ok("DescNullsLast"), + Ok(SortDirection::Clustered) => Ok("Clustered"), + Ok(SortDirection::Unspecified) => Err(PlanError::invalid( + "SortField", + Some("sort_kind"), + "Unspecified SortDirection", + )), + Err(_) => Err(PlanError::invalid( + "SortField", + Some("sort_kind"), + format!("Unknown SortDirection: {d}"), + )), + }, + Some(SortKind::ComparisonFunctionReference(f)) => Err(PlanError::unimplemented( + "SortField", + Some("sort_kind"), + format!("ComparisonFunctionReference {f} textification not implemented"), + )), + None => Err(PlanError::invalid( + "SortField", + Some("sort_kind"), + "Missing sort_kind", + )), + } +} + +/// Write a single sort field as `($ref,&Direction)`, for use in `order=`. +fn textify_sort_field(sf: &SortField, ctx: &S, w: &mut W) -> fmt::Result { + let expr = ctx.expect(sf.expr.as_ref()); + write!(w, "({expr},")?; + match sort_direction_str(sf.sort_kind.as_ref()) { + Ok(s) => textify_enum(s, ctx, w)?, + Err(e) => write!(w, "{}", ctx.failure(e))?, + } + write!(w, ")") +} + +/// Resolve an `AggregationPhase` value to its enum variant name, for `phase=`. +fn aggregation_phase_str(phase: i32) -> Result<&'static str, PlanError> { + match AggregationPhase::try_from(phase) { + Ok(AggregationPhase::Unspecified) => Ok("Unspecified"), + Ok(AggregationPhase::InitialToIntermediate) => Ok("InitialToIntermediate"), + Ok(AggregationPhase::IntermediateToIntermediate) => Ok("IntermediateToIntermediate"), + Ok(AggregationPhase::InitialToResult) => Ok("InitialToResult"), + Ok(AggregationPhase::IntermediateToResult) => Ok("IntermediateToResult"), + Err(_) => Err(PlanError::invalid( + "WindowFunction", + Some("phase"), + format!("Unknown AggregationPhase: {phase}"), + )), + } +} + +/// Resolve an `AggregationInvocation` value to its enum variant name, for `invocation=`. +fn aggregation_invocation_str(invocation: i32) -> Result<&'static str, PlanError> { + match AggregationInvocation::try_from(invocation) { + Ok(AggregationInvocation::Unspecified) => Ok("Unspecified"), + Ok(AggregationInvocation::All) => Ok("All"), + Ok(AggregationInvocation::Distinct) => Ok("Distinct"), + Err(_) => Err(PlanError::invalid( + "WindowFunction", + Some("invocation"), + format!("Unknown AggregationInvocation: {invocation}"), + )), + } +} + +/// Write a window bound as an integer offset (negative preceding, positive +/// following, `0` for the current row) or `_` for unbounded/unspecified. +fn textify_window_bound( + bound: Option<&window_function::Bound>, + _ctx: &S, + w: &mut W, +) -> fmt::Result { + match bound.and_then(|b| b.kind.as_ref()) { + None | Some(bound::Kind::Unbounded(_)) => write!(w, "_"), + Some(bound::Kind::CurrentRow(_)) => write!(w, "0"), + Some(bound::Kind::Preceding(p)) => write!(w, "{}", -p.offset), + Some(bound::Kind::Following(f)) => write!(w, "{}", f.offset), + } +} + +impl Textify for WindowFunction { + fn name() -> &'static str { + "WindowFunction" + } + + fn textify(&self, ctx: &S, w: &mut W) -> fmt::Result { + // Shared with ScalarFunction/AggregateFunction textification + let name_and_anchor = + NamedAnchor::lookup(ctx, ExtensionKind::Function, self.function_reference); + let name_and_anchor = ctx.display(&name_and_anchor); + + let args = ctx.separated(&self.arguments, ", "); + let options = ctx.separated(&self.options, ", "); + let between = if self.arguments.is_empty() || self.options.is_empty() { + "" + } else { + ", " + }; + + let output = OutputType(self.output_type.as_ref()); + let output_type = ctx.display(&output); + + write!(w, "{name_and_anchor}({args}{between}{options}) over(")?; + + let mut named_args: Vec = Vec::new(); + + let mut phase_str = String::new(); + match aggregation_phase_str(self.phase) { + Ok(s) => textify_enum(s, ctx, &mut phase_str)?, + Err(e) => write!(phase_str, "{}", ctx.failure(e))?, + } + named_args.push(format!("phase={phase_str}")); + + // order= is omitted when there are no sort fields. A single sort field + // is written bare; two or more are wrapped in a parenthesized list of + // tuples. + if !self.sorts.is_empty() { + let mut order_str = String::new(); + if self.sorts.len() == 1 { + textify_sort_field(&self.sorts[0], ctx, &mut order_str)?; + } else { + write!(order_str, "(")?; + for (i, sf) in self.sorts.iter().enumerate() { + if i > 0 { + write!(order_str, ",")?; + } + textify_sort_field(sf, ctx, &mut order_str)?; + } + write!(order_str, ")")?; + } + named_args.push(format!("order={order_str}")); + } + + if self.invocation != AggregationInvocation::Unspecified as i32 { + let mut invocation_str = String::new(); + match aggregation_invocation_str(self.invocation) { + Ok(s) => textify_enum(s, ctx, &mut invocation_str)?, + Err(e) => write!(invocation_str, "{}", ctx.failure(e))?, + } + named_args.push(format!("invocation={invocation_str}")); + } + + if !self.partitions.is_empty() { + let parts = ctx.separated(&self.partitions, ", "); + named_args.push(format!("partition=({parts})")); + } + + let bounds_type = window_function::BoundsType::try_from(self.bounds_type); + let has_bounds = self.lower_bound.is_some() || self.upper_bound.is_some(); + if !matches!(bounds_type, Ok(window_function::BoundsType::Unspecified)) || has_bounds { + let keyword = match bounds_type { + Ok(window_function::BoundsType::Rows) => "rows", + Ok(window_function::BoundsType::Range) => "range", + Ok(window_function::BoundsType::Unspecified) => { + // bounds_type is required whenever a frame is present. + ctx.push_error( + PlanError::invalid( + "WindowFunction", + Some("bounds_type"), + "bounds_type is Unspecified but lower_bound/upper_bound are set", + ) + .into(), + ); + "rows" + } + Err(_) => { + ctx.push_error( + PlanError::invalid( + "WindowFunction", + Some("bounds_type"), + format!("Unknown BoundsType: {}", self.bounds_type), + ) + .into(), + ); + "rows" + } + }; + let mut lower_str = String::new(); + textify_window_bound(self.lower_bound.as_ref(), ctx, &mut lower_str)?; + let mut upper_str = String::new(); + textify_window_bound(self.upper_bound.as_ref(), ctx, &mut upper_str)?; + named_args.push(format!("{keyword}=({lower_str}, {upper_str})")); + } + + write!(w, "{}", named_args.join(", "))?; + write!(w, "){output_type}") + } +} + impl Textify for RexType { fn name() -> &'static str { "RexType" @@ -532,15 +735,7 @@ impl Textify for RexType { RexType::Literal(literal) => literal.textify(ctx, w), RexType::Selection(f) => f.textify(ctx, w), RexType::ScalarFunction(s) => s.textify(ctx, w), - RexType::WindowFunction(_w) => write!( - w, - "{}", - ctx.failure(PlanError::unimplemented( - "RexType", - Some("WindowFunction"), - "WindowFunction textification not implemented", - )) - ), + RexType::WindowFunction(f) => f.textify(ctx, w), RexType::IfThen(i) => i.textify(ctx, w), RexType::SwitchExpression(_s) => write!( w, @@ -1135,4 +1330,119 @@ mod tests { }; assert_eq!(ctx.textify_no_errors(&cast), "(1:i32)::json"); } + + fn base_window_function() -> WindowFunction { + WindowFunction { + function_reference: 10, + arguments: vec![], + options: vec![], + output_type: Some(make_i16_type()), + phase: AggregationPhase::InitialToResult as i32, + sorts: vec![], + invocation: AggregationInvocation::Unspecified as i32, + partitions: vec![], + bounds_type: window_function::BoundsType::Unspecified as i32, + lower_bound: None, + upper_bound: None, + #[allow(deprecated)] + args: vec![], + } + } + + fn field_expr(field: i32) -> Expression { + Expression { + rex_type: Some(RexType::Selection(Box::new(struct_field_reference(field)))), + } + } + + fn sort_field(field: i32, direction: SortDirection) -> SortField { + SortField { + expr: Some(field_expr(field)), + sort_kind: Some(SortKind::Direction(direction as i32)), + } + } + + #[test] + fn test_window_function_no_bound() { + let ctx = TestContext::new() + .with_urn(1, "urn:example") + .with_function(1, 10, "row_number"); + let f = base_window_function(); + assert_eq!( + ctx.textify_no_errors(&f), + "row_number() over(phase=&InitialToResult):i16" + ); + } + + #[test] + fn test_window_function_full() { + let ctx = TestContext::new() + .with_urn(1, "urn:example") + .with_function(1, 10, "sum"); + let mut f = base_window_function(); + f.partitions = vec![field_expr(0)]; + f.sorts = vec![sort_field(1, SortDirection::AscNullsLast)]; + f.invocation = AggregationInvocation::Distinct as i32; + f.bounds_type = window_function::BoundsType::Rows as i32; + f.lower_bound = Some(window_function::Bound { + kind: Some(bound::Kind::Preceding(bound::Preceding { offset: 3 })), + }); + f.upper_bound = Some(window_function::Bound { + kind: Some(bound::Kind::CurrentRow(bound::CurrentRow {})), + }); + assert_eq!( + ctx.textify_no_errors(&f), + "sum() over(phase=&InitialToResult, order=($1,&AscNullsLast), invocation=&Distinct, partition=($0), rows=(-3, 0)):i16" + ); + } + + #[test] + fn test_window_function_multiple_order_fields() { + let ctx = TestContext::new() + .with_urn(1, "urn:example") + .with_function(1, 10, "sum"); + let mut f = base_window_function(); + f.sorts = vec![ + sort_field(0, SortDirection::AscNullsLast), + sort_field(1, SortDirection::DescNullsFirst), + ]; + assert_eq!( + ctx.textify_no_errors(&f), + "sum() over(phase=&InitialToResult, order=(($0,&AscNullsLast),($1,&DescNullsFirst))):i16" + ); + } + + #[test] + fn test_window_function_unbounded_bounds() { + let ctx = TestContext::new() + .with_urn(1, "urn:example") + .with_function(1, 10, "sum"); + let mut f = base_window_function(); + f.bounds_type = window_function::BoundsType::Rows as i32; + f.lower_bound = None; + f.upper_bound = Some(window_function::Bound { + kind: Some(bound::Kind::Unbounded(bound::Unbounded {})), + }); + assert_eq!( + ctx.textify_no_errors(&f), + "sum() over(phase=&InitialToResult, rows=(_, _)):i16" + ); + } + + #[test] + fn test_window_function_unspecified_bounds_type_with_bounds_is_lossy() { + // bounds_type is Unspecified, but a frame is set anyway. + // The textifier still emits best-effort rows=/range= output but surfaces + // the loss as an accumulated error rather than silently dropping it. + let ctx = TestContext::new() + .with_urn(1, "urn:example") + .with_function(1, 10, "sum"); + let mut f = base_window_function(); + f.lower_bound = Some(window_function::Bound { + kind: Some(bound::Kind::CurrentRow(bound::CurrentRow {})), + }); + let (s, errs) = ctx.textify(&f); + assert_eq!(s, "sum() over(phase=&InitialToResult, rows=(0, _)):i16"); + assert!(!errs.is_empty(), "expected a diagnostic about bounds_type"); + } } diff --git a/tests/plan_roundtrip.rs b/tests/plan_roundtrip.rs index bd85454e..6bd56711 100644 --- a/tests/plan_roundtrip.rs +++ b/tests/plan_roundtrip.rs @@ -1133,3 +1133,89 @@ Root[a] assert!(Parser::parse(plan).is_err()); } + +#[test] +fn test_window_function_all_named_args_roundtrip() { + let plan = r#"=== Extensions +URNs: + @ 1: https://github.com/substrait-io/substrait/blob/main/extensions/functions_aggregate.yaml +Functions: + # 10 @ 1: sum + +=== Plan +Root[a, b, s] + Project[$0, $1, sum($1) over(phase=&InitialToResult, order=($1,&AscNullsLast), invocation=&Distinct, partition=($0), rows=(-3, 0)):fp64?] + Read[t => a:i32, b:fp64]"#; + + roundtrip_plan(plan); +} + +#[test] +fn test_window_function_unbounded_roundtrip() { + let plan = r#"=== Extensions +URNs: + @ 1: https://github.com/substrait-io/substrait/blob/main/extensions/functions_arithmetic.yaml +Functions: + # 10 @ 1: row_number + +=== Plan +Root[a, r] + Project[$0, row_number() over(phase=&InitialToResult, order=($0,&AscNullsLast), partition=($0)):i64] + Read[t => a:i32]"#; + + roundtrip_plan(plan); +} + +/// Minimal, isolated regression test for the bare (non-tuple) rendering of a +/// single `order=` sort field. +#[test] +fn test_window_function_single_order_field_bare_form() { + let plan = r#"=== Extensions +URNs: + @ 1: https://github.com/substrait-io/substrait/blob/main/extensions/functions_arithmetic.yaml +Functions: + # 10 @ 1: row_number + +=== Plan +Root[r] + Project[row_number() over(phase=&InitialToResult, order=($0,&AscNullsLast)):i64] + Read[t => a:i32]"#; + + roundtrip_plan(plan); +} + +/// `_` for an unbounded lower bound round-trips as `_`, alongside a bounded +/// (current row) upper bound. +#[test] +fn test_window_function_unbounded_lower_bound_roundtrip() { + let plan = r#"=== Extensions +URNs: + @ 1: https://github.com/substrait-io/substrait/blob/main/extensions/functions_arithmetic.yaml +Functions: + # 10 @ 1: sum + +=== Plan +Root[s] + Project[sum($0) over(phase=&InitialToResult, rows=(_, 0)):i64?] + Read[t => a:i64]"#; + + roundtrip_plan(plan); +} + +/// A `range=` frame with more than one `order=` field is rejected at parse +/// time: Substrait only allows `RANGE` bounds with a single ordering column. +#[test] +fn test_window_function_range_frame_requires_single_order_field() { + let plan = r#"=== Extensions +URNs: + @ 1: https://github.com/substrait-io/substrait/blob/main/extensions/functions_arithmetic.yaml +Functions: + # 10 @ 1: sum + +=== Plan +Root[s] + Project[sum($0) over(order=(($0,&AscNullsLast),($1,&AscNullsLast)), range=(_, 0), phase=&InitialToResult):i64?] + Read[t => a:i64, b:i64]"#; + + assert!(Parser::parse(plan).is_err()); +} From 8f7dd9068af20f7d20daadb7e43ac3bb00f7fbde Mon Sep 17 00:00:00 2001 From: "gordon.hamilton" Date: Wed, 15 Jul 2026 20:35:47 -0400 Subject: [PATCH 2/7] feat: depulicating and cleaning up code based on review --- GRAMMAR.md | 6 +- src/parser/expression_grammar.pest | 8 +- src/parser/expressions.rs | 213 ++++++++++++++------------- src/textify/expressions.rs | 226 +++++++++++++---------------- src/textify/rels.rs | 32 +++- tests/plan_roundtrip.rs | 36 +++++ 6 files changed, 292 insertions(+), 229 deletions(-) diff --git a/GRAMMAR.md b/GRAMMAR.md index 0fa62465..5c147311 100644 --- a/GRAMMAR.md +++ b/GRAMMAR.md @@ -507,7 +507,7 @@ A window function computes a value over a "window" of rows related to the curren #### Syntax -`window_function := function_signature anchor? urn_anchor? "(" (expression ("," expression)*)? ")" "over(" window_named_arg ("," window_named_arg)* ")" ":" type` +`window_function := function_signature anchor? urn_anchor? "(" (expression ("," expression)*)? ")" "over(" (window_named_arg ("," window_named_arg)*)? ")" ":" type` `window_named_arg` is one of: @@ -518,9 +518,7 @@ A window function computes a value over a "window" of rows related to the curren - `rows=(lower, upper)` / `range=(lower, upper)` - the window frame, each bound is an integer (negative = preceding, positive = following, `0` = current row) or `_` for unbounded -`partition=`, `order=`, `invocation=`, and `rows=`/`range=` are all -optional and are omitted when empty; `phase=` is always required and -always printed. A `range=` frame requires exactly one `order=` field. +`partition=`, `order=`, `invocation=`, and `rows=`/`range=` are all optional and are omitted when empty; `phase=` is always required and always printed, though this is enforced when parsing rather than by the grammar itself (the named-arg list as a whole is syntactically optional). A `range=` frame requires exactly one `order=` field. #### Examples diff --git a/src/parser/expression_grammar.pest b/src/parser/expression_grammar.pest index 66135b03..5562ff81 100644 --- a/src/parser/expression_grammar.pest +++ b/src/parser/expression_grammar.pest @@ -153,8 +153,12 @@ argument_list = { "(" ~ (expression ~ (sp ~ "," ~ sp ~ expression)*)? ~ ")" } // - Arguments `()` are required, e.g. (1, 2, 3) // - Required output type annotation after closing paren, e.g. :i64 // (unambiguous because function_signature is atomic and ends before the opening paren) +function_call_prefix = _{ + function_signature ~ sp ~ anchor? ~ sp ~ urn_anchor? ~ sp ~ argument_list +} + function_call = { - function_signature ~ sp ~ anchor? ~ sp ~ urn_anchor? ~ sp ~ argument_list ~ ":" ~ sp ~ type + function_call_prefix ~ ":" ~ sp ~ type } if_clause = { @@ -207,7 +211,7 @@ window_named_arg = { window_partition_arg | window_order_arg | window_frame window_named_arg_list = { (window_named_arg ~ (sp ~ "," ~ sp ~ window_named_arg)*)? } window_function_call = { - function_signature ~ sp ~ anchor? ~ sp ~ urn_anchor? ~ sp ~ argument_list ~ sp ~ "over" ~ sp ~ "(" ~ sp ~ window_named_arg_list ~ sp ~ ")" ~ ":" ~ sp ~ type + function_call_prefix ~ sp ~ "over" ~ sp ~ "(" ~ sp ~ window_named_arg_list ~ sp ~ ")" ~ ":" ~ sp ~ type } // Top-level Expression Rule diff --git a/src/parser/expressions.rs b/src/parser/expressions.rs index fafc48e8..23e15cf7 100644 --- a/src/parser/expressions.rs +++ b/src/parser/expressions.rs @@ -22,6 +22,7 @@ use super::{ }; use crate::extensions::SimpleExtensions; use crate::extensions::simple::{CompoundName, ExtensionKind}; +use crate::parser::relations::parse_expression_list; /// A field index (e.g., parsed from "$0" -> 0). #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -400,6 +401,47 @@ impl ScopedParsePair for Literal { } } +/// The shared prefix of a scalar/window function call. +struct FunctionHead { + name: CompoundName, + anchor: Option, + arguments: Vec, +} + +/// Parse the `CompoundName ~ anchor? ~ urn_anchor? ~ argument_list` prefix shared by `ScalarFunction` and `WindowFunction`. +fn parse_function_head( + extensions: &SimpleExtensions, + iter: &mut RuleIter<'_>, +) -> Result { + // Parse compound function name (required) — e.g. "equal" or "equal:any_any" + let name = iter.parse_next::(); + + // Parse optional anchor (e.g., #1) + let anchor = iter + .try_pop(Rule::anchor) + .map(|n| unwrap_single_pair(n).as_str().parse::().unwrap()); + + // Parse optional URN anchor (e.g., @1) + let _urn_anchor = iter + .try_pop(Rule::urn_anchor) + .map(|n| unwrap_single_pair(n).as_str().parse::().unwrap()); + + // Parse argument list (required) + let argument_list = iter.pop(Rule::argument_list); + let mut arguments = Vec::new(); + for e in argument_list.into_inner() { + arguments.push(FunctionArgument { + arg_type: Some(ArgType::Value(Expression::parse_pair(extensions, e)?)), + }); + } + + Ok(FunctionHead { + name, + anchor, + arguments, + }) +} + impl ScopedParsePair for ScalarFunction { fn rule() -> Rule { Rule::function_call @@ -417,27 +459,11 @@ impl ScopedParsePair for ScalarFunction { let span = pair.as_span(); let mut iter = RuleIter::from(pair.into_inner()); - // Parse compound function name (required) — e.g. "equal" or "equal:any_any" - let name = iter.parse_next::(); - - // Parse optional anchor (e.g., #1) - let anchor = iter - .try_pop(Rule::anchor) - .map(|n| unwrap_single_pair(n).as_str().parse::().unwrap()); - - // Parse optional URN anchor (e.g., @1) - let _urn_anchor = iter - .try_pop(Rule::urn_anchor) - .map(|n| unwrap_single_pair(n).as_str().parse::().unwrap()); - - // Parse argument list (required) - let argument_list = iter.pop(Rule::argument_list); - let mut arguments = Vec::new(); - for e in argument_list.into_inner() { - arguments.push(FunctionArgument { - arg_type: Some(ArgType::Value(Expression::parse_pair(extensions, e)?)), - }); - } + let FunctionHead { + name, + anchor, + arguments, + } = parse_function_head(extensions, &mut iter)?; // Parse required output type (e.g., :i64). pop is safe here because // the grammar guarantees the type token is always present. @@ -486,48 +512,51 @@ fn parse_window_bound(pair: pest::iterators::Pair) -> Option) -> Result { - assert_eq!(pair.as_rule(), Rule::enum_value); - let name = pair.as_str().trim_start_matches('&'); - let phase = match name { - "Unspecified" => AggregationPhase::Unspecified, - "InitialToIntermediate" => AggregationPhase::InitialToIntermediate, - "IntermediateToIntermediate" => AggregationPhase::IntermediateToIntermediate, - "InitialToResult" => AggregationPhase::InitialToResult, - "IntermediateToResult" => AggregationPhase::IntermediateToResult, - other => { - return Err(MessageParseError::invalid( - "AggregationPhase", - pair.as_span(), - format!("Unknown aggregation phase: {other}"), - )); - } - }; - Ok(phase as i32) -} - -/// Parse an `enum_value` pair (`&Identifier`) into an `AggregationInvocation`. -fn parse_aggregation_invocation( +/// Parse an `enum_value` pair (`&Identifier`) into the `i32` discriminant of +/// one of `variants` (name, discriminant pairs), for enums that are looked up by name. +fn parse_enum_value_by_name( pair: pest::iterators::Pair, + type_name: &'static str, + variants: &[(&str, i32)], ) -> Result { assert_eq!(pair.as_rule(), Rule::enum_value); let name = pair.as_str().trim_start_matches('&'); - let invocation = match name { - "Unspecified" => AggregationInvocation::Unspecified, - "All" => AggregationInvocation::All, - "Distinct" => AggregationInvocation::Distinct, - other => { - return Err(MessageParseError::invalid( - "AggregationInvocation", + variants + .iter() + .find(|(variant, _)| *variant == name) + .map(|(_, value)| *value) + .ok_or_else(|| { + MessageParseError::invalid( + type_name, pair.as_span(), - format!("Unknown aggregation invocation: {other}"), - )); - } - }; - Ok(invocation as i32) + format!("Unknown {type_name}: {name}"), + ) + }) } +const AGGREGATION_PHASE_VARIANTS: &[(&str, i32)] = &[ + ("Unspecified", AggregationPhase::Unspecified as i32), + ( + "InitialToIntermediate", + AggregationPhase::InitialToIntermediate as i32, + ), + ( + "IntermediateToIntermediate", + AggregationPhase::IntermediateToIntermediate as i32, + ), + ("InitialToResult", AggregationPhase::InitialToResult as i32), + ( + "IntermediateToResult", + AggregationPhase::IntermediateToResult as i32, + ), +]; + +const AGGREGATION_INVOCATION_VARIANTS: &[(&str, i32)] = &[ + ("Unspecified", AggregationInvocation::Unspecified as i32), + ("All", AggregationInvocation::All as i32), + ("Distinct", AggregationInvocation::Distinct as i32), +]; + impl ScopedParsePair for WindowFunction { fn rule() -> Rule { Rule::window_function_call @@ -545,27 +574,11 @@ impl ScopedParsePair for WindowFunction { let span = pair.as_span(); let mut iter = RuleIter::from(pair.into_inner()); - // Parse compound function name (required) — e.g. "row_number" or "sum:i64" - let name = iter.parse_next::(); - - // Parse optional anchor (e.g., #1) - let anchor = iter - .try_pop(Rule::anchor) - .map(|n| unwrap_single_pair(n).as_str().parse::().unwrap()); - - // Parse optional URN anchor (e.g., @1) - let _urn_anchor = iter - .try_pop(Rule::urn_anchor) - .map(|n| unwrap_single_pair(n).as_str().parse::().unwrap()); - - // Parse argument list (required) - let argument_list = iter.pop(Rule::argument_list); - let mut arguments = Vec::new(); - for e in argument_list.into_inner() { - arguments.push(FunctionArgument { - arg_type: Some(ArgType::Value(Expression::parse_pair(extensions, e)?)), - }); - } + let FunctionHead { + name, + anchor, + arguments, + } = parse_function_head(extensions, &mut iter)?; // Parse the required `over(...)` named-argument list let named_arg_list = iter.pop(Rule::window_named_arg_list); @@ -589,9 +602,7 @@ impl ScopedParsePair for WindowFunction { Rule::window_partition_arg => { let mut parts_iter = RuleIter::from(inner.into_inner()); if let Some(expr_list) = parts_iter.try_pop(Rule::expression_list) { - for e in expr_list.into_inner() { - partitions.push(Expression::parse_pair(extensions, e)?); - } + partitions = parse_expression_list(extensions, expr_list)?; } parts_iter.done(); } @@ -623,13 +634,21 @@ impl ScopedParsePair for WindowFunction { let mut phase_iter = RuleIter::from(inner.into_inner()); let enum_pair = phase_iter.pop(Rule::enum_value); phase_iter.done(); - phase = Some(parse_aggregation_phase(enum_pair)?); + phase = Some(parse_enum_value_by_name( + enum_pair, + "AggregationPhase", + AGGREGATION_PHASE_VARIANTS, + )?); } Rule::window_invocation_arg => { let mut inv_iter = RuleIter::from(inner.into_inner()); let enum_pair = inv_iter.pop(Rule::enum_value); inv_iter.done(); - invocation = parse_aggregation_invocation(enum_pair)?; + invocation = parse_enum_value_by_name( + enum_pair, + "AggregationInvocation", + AGGREGATION_INVOCATION_VARIANTS, + )?; } other => { unreachable!("Grammar guarantees window_named_arg alternatives, got {other:?}") @@ -1384,9 +1403,7 @@ mod tests { assert_eq!(pairs.as_str(), "equal:any_any"); } - // ---- Tests for ScalarFunction parsing with compound names ---- - - fn make_extensions_for_fn_tests() -> SimpleExtensions { + fn make_extensions() -> SimpleExtensions { let mut exts = SimpleExtensions::default(); exts.add_extension_urn("urn".to_string(), 1).unwrap(); exts.add_extension( @@ -1416,7 +1433,7 @@ mod tests { #[test] fn test_scalar_function_full_compound_name() { // Full compound name without anchor - let exts = make_extensions_for_fn_tests(); + let exts = make_extensions(); let pair = parse_exact(Rule::function_call, "equal:any_any($0, $1):boolean"); let f = ScalarFunction::parse_pair(&exts, pair).unwrap(); assert_eq!(f.function_reference, 1); @@ -1429,7 +1446,7 @@ mod tests { #[test] fn test_scalar_function_second_overload() { - let exts = make_extensions_for_fn_tests(); + let exts = make_extensions(); let pair = parse_exact(Rule::function_call, "equal:str_str($0, $1):boolean"); let f = ScalarFunction::parse_pair(&exts, pair).unwrap(); @@ -1440,7 +1457,7 @@ mod tests { #[test] fn test_scalar_function_base_name_unique_overload() { // "add" has only one overload; base-name lookup should succeed - let exts = make_extensions_for_fn_tests(); + let exts = make_extensions(); let pair = parse_exact(Rule::function_call, "add($0, $1):i64"); let f = ScalarFunction::parse_pair(&exts, pair).unwrap(); @@ -1455,7 +1472,7 @@ mod tests { #[test] fn test_scalar_function_base_name_ambiguous_fails() { // "equal" has two overloads; base-name lookup should fail - let exts = make_extensions_for_fn_tests(); + let exts = make_extensions(); let pair = parse_exact(Rule::function_call, "equal($0, $1):boolean"); let result = ScalarFunction::parse_pair(&exts, pair); assert!(result.is_err(), "ambiguous base name should fail"); @@ -1463,7 +1480,7 @@ mod tests { #[test] fn test_scalar_function_compound_name_with_anchor() { - let exts = make_extensions_for_fn_tests(); + let exts = make_extensions(); let pair = parse_exact(Rule::function_call, "equal:any_any#1($0, $1):boolean"); let f = ScalarFunction::parse_pair(&exts, pair).unwrap(); assert_eq!(f.function_reference, 1); @@ -1473,7 +1490,7 @@ mod tests { #[test] fn test_scalar_function_base_name_with_anchor() { // Base name + explicit anchor should resolve (anchor 1 stores equal:any_any) - let exts = make_extensions_for_fn_tests(); + let exts = make_extensions(); let pair = parse_exact(Rule::function_call, "equal#1($0, $1):boolean"); let f = ScalarFunction::parse_pair(&exts, pair).unwrap(); assert_eq!(f.function_reference, 1); @@ -1482,7 +1499,7 @@ mod tests { #[test] fn test_scalar_function_wrong_name_for_anchor_fails() { - let exts = make_extensions_for_fn_tests(); + let exts = make_extensions(); let pair = parse_exact(Rule::function_call, "like#1($0):boolean"); let result = ScalarFunction::parse_pair(&exts, pair); assert!(result.is_err(), "mismatched name/anchor should fail"); @@ -1681,7 +1698,7 @@ mod tests { #[test] fn test_window_function_full() { - let exts = make_extensions_for_fn_tests(); + let exts = make_extensions(); let pair = parse_exact( Rule::window_function_call, "add:i64_i64($0, $1) over(partition=($0), order=($1,&AscNullsLast), invocation=&Distinct, rows=(-3, 0), phase=&InitialToResult):i64", @@ -1710,7 +1727,7 @@ mod tests { #[test] fn test_window_function_unbounded() { - let exts = make_extensions_for_fn_tests(); + let exts = make_extensions(); let pair = parse_exact( Rule::window_function_call, "add:i64_i64($0, $1) over(phase=&InitialToResult):i64", @@ -1729,7 +1746,7 @@ mod tests { #[test] fn test_window_function_unbounded_lower_bound() { - let exts = make_extensions_for_fn_tests(); + let exts = make_extensions(); let pair = parse_exact( Rule::window_function_call, "add:i64_i64($0, $1) over(rows=(_, 5), phase=&InitialToResult):i64", @@ -1751,7 +1768,7 @@ mod tests { #[test] fn test_window_function_missing_phase_fails() { - let exts = make_extensions_for_fn_tests(); + let exts = make_extensions(); let pair = parse_exact( Rule::window_function_call, "add:i64_i64($0, $1) over(partition=($0)):i64", @@ -1762,7 +1779,7 @@ mod tests { #[test] fn test_window_function_range_with_multiple_order_fields_fails() { - let exts = make_extensions_for_fn_tests(); + let exts = make_extensions(); let pair = parse_exact( Rule::window_function_call, "add:i64_i64($0, $1) over(order=(($0,&AscNullsLast),($1,&AscNullsLast)), range=(_, 0), phase=&InitialToResult):i64", @@ -1776,7 +1793,7 @@ mod tests { #[test] fn test_window_function_call_in_expression() { - let exts = make_extensions_for_fn_tests(); + let exts = make_extensions(); let pair = parse_exact( Rule::expression, "add:i64_i64($0, $1) over(phase=&InitialToResult):i64", diff --git a/src/textify/expressions.rs b/src/textify/expressions.rs index 8a738956..a7491f17 100644 --- a/src/textify/expressions.rs +++ b/src/textify/expressions.rs @@ -11,7 +11,6 @@ use substrait::proto::expression::{ reference_segment, }; use substrait::proto::function_argument::ArgType; -use substrait::proto::sort_field::{SortDirection, SortKind}; use substrait::proto::{ AggregateFunction, AggregationPhase, Expression, FunctionArgument, FunctionOption, SortField, expression as expr, @@ -19,6 +18,7 @@ use substrait::proto::{ use super::{PlanError, Scope, Textify, Visibility}; use crate::extensions::simple::ExtensionKind; +use crate::textify::rels::ValueEnum; use crate::textify::types::{Name, NamedAnchor, OutputType, escaped}; // …(…) for function call @@ -407,32 +407,45 @@ impl Textify for FieldReference { } } +/// Write the `name(args, options)` prefix shared by `ScalarFunction`, +/// `AggregateFunction`, and `WindowFunction` textification. +fn textify_function_call_prefix( + function_reference: u32, + arguments: &[FunctionArgument], + options: &[FunctionOption], + ctx: &S, + w: &mut W, +) -> fmt::Result { + let name_and_anchor = NamedAnchor::lookup(ctx, ExtensionKind::Function, function_reference); + let name_and_anchor = ctx.display(&name_and_anchor); + + let between = if arguments.is_empty() || options.is_empty() { + "" + } else { + ", " + }; + let args = ctx.separated(arguments, ", "); + let options = ctx.separated(options, ", "); + + write!(w, "{name_and_anchor}({args}{between}{options})") +} + impl Textify for ScalarFunction { fn name() -> &'static str { "ScalarFunction" } fn textify(&self, ctx: &S, w: &mut W) -> fmt::Result { - let name_and_anchor = - NamedAnchor::lookup(ctx, ExtensionKind::Function, self.function_reference); - let name_and_anchor = ctx.display(&name_and_anchor); - - let args = ctx.separated(&self.arguments, ", "); - let options = ctx.separated(&self.options, ", "); - let between = if self.arguments.is_empty() || self.options.is_empty() { - "" - } else { - ", " - }; - - let output = OutputType(self.output_type.as_ref()); - let output_type = ctx.display(&output); - - write!( + textify_function_call_prefix( + self.function_reference, + &self.arguments, + &self.options, + ctx, w, - "{name_and_anchor}({args}{between}{options}){output_type}" )?; - Ok(()) + + let output = OutputType(self.output_type.as_ref()); + write!(w, "{}", ctx.display(&output)) } } @@ -527,85 +540,50 @@ impl Textify for IfThen { } } -/// Resolve a `SortField`'s direction to its enum variant name, for `order=`. -fn sort_direction_str(sk: Option<&SortKind>) -> Result<&'static str, PlanError> { - match sk { - Some(SortKind::Direction(d)) => match SortDirection::try_from(*d) { - Ok(SortDirection::AscNullsFirst) => Ok("AscNullsFirst"), - Ok(SortDirection::AscNullsLast) => Ok("AscNullsLast"), - Ok(SortDirection::DescNullsFirst) => Ok("DescNullsFirst"), - Ok(SortDirection::DescNullsLast) => Ok("DescNullsLast"), - Ok(SortDirection::Clustered) => Ok("Clustered"), - Ok(SortDirection::Unspecified) => Err(PlanError::invalid( - "SortField", - Some("sort_kind"), - "Unspecified SortDirection", - )), - Err(_) => Err(PlanError::invalid( - "SortField", - Some("sort_kind"), - format!("Unknown SortDirection: {d}"), - )), - }, - Some(SortKind::ComparisonFunctionReference(f)) => Err(PlanError::unimplemented( - "SortField", - Some("sort_kind"), - format!("ComparisonFunctionReference {f} textification not implemented"), - )), +/// Write a single sort field as `($ref,&Direction)`, for use in `order=`. +fn textify_sort_field(sf: &SortField, ctx: &S, w: &mut W) -> fmt::Result { + let expr = ctx.expect(sf.expr.as_ref()); + write!(w, "({expr},")?; + let result = match sf.sort_kind.as_ref() { + Some(sk) => sk.as_enum_str(), None => Err(PlanError::invalid( "SortField", Some("sort_kind"), "Missing sort_kind", )), - } -} - -/// Write a single sort field as `($ref,&Direction)`, for use in `order=`. -fn textify_sort_field(sf: &SortField, ctx: &S, w: &mut W) -> fmt::Result { - let expr = ctx.expect(sf.expr.as_ref()); - write!(w, "({expr},")?; - match sort_direction_str(sf.sort_kind.as_ref()) { - Ok(s) => textify_enum(s, ctx, w)?, + }; + match result { + Ok(s) => textify_enum(&s, ctx, w)?, Err(e) => write!(w, "{}", ctx.failure(e))?, } write!(w, ")") } -/// Resolve an `AggregationPhase` value to its enum variant name, for `phase=`. -fn aggregation_phase_str(phase: i32) -> Result<&'static str, PlanError> { - match AggregationPhase::try_from(phase) { - Ok(AggregationPhase::Unspecified) => Ok("Unspecified"), - Ok(AggregationPhase::InitialToIntermediate) => Ok("InitialToIntermediate"), - Ok(AggregationPhase::IntermediateToIntermediate) => Ok("IntermediateToIntermediate"), - Ok(AggregationPhase::InitialToResult) => Ok("InitialToResult"), - Ok(AggregationPhase::IntermediateToResult) => Ok("IntermediateToResult"), +/// Resolve a raw `i32` discriminant to `T`, then render its name via `ValueEnum`, for enum fields. +fn textify_i32_enum(raw: i32, type_name: &'static str, ctx: &S, w: &mut W) -> fmt::Result +where + T: TryFrom + ValueEnum, + S: Scope, + W: fmt::Write, +{ + let result = match T::try_from(raw) { + Ok(v) => v.as_enum_str(), Err(_) => Err(PlanError::invalid( "WindowFunction", - Some("phase"), - format!("Unknown AggregationPhase: {phase}"), - )), - } -} - -/// Resolve an `AggregationInvocation` value to its enum variant name, for `invocation=`. -fn aggregation_invocation_str(invocation: i32) -> Result<&'static str, PlanError> { - match AggregationInvocation::try_from(invocation) { - Ok(AggregationInvocation::Unspecified) => Ok("Unspecified"), - Ok(AggregationInvocation::All) => Ok("All"), - Ok(AggregationInvocation::Distinct) => Ok("Distinct"), - Err(_) => Err(PlanError::invalid( - "WindowFunction", - Some("invocation"), - format!("Unknown AggregationInvocation: {invocation}"), + Some(type_name), + format!("Unknown {type_name}: {raw}"), )), + }; + match result { + Ok(s) => textify_enum(&s, ctx, w), + Err(e) => write!(w, "{}", ctx.failure(e)), } } /// Write a window bound as an integer offset (negative preceding, positive /// following, `0` for the current row) or `_` for unbounded/unspecified. -fn textify_window_bound( +fn textify_window_bound( bound: Option<&window_function::Bound>, - _ctx: &S, w: &mut W, ) -> fmt::Result { match bound.and_then(|b| b.kind.as_ref()) { @@ -622,31 +600,24 @@ impl Textify for WindowFunction { } fn textify(&self, ctx: &S, w: &mut W) -> fmt::Result { - // Shared with ScalarFunction/AggregateFunction textification - let name_and_anchor = - NamedAnchor::lookup(ctx, ExtensionKind::Function, self.function_reference); - let name_and_anchor = ctx.display(&name_and_anchor); - - let args = ctx.separated(&self.arguments, ", "); - let options = ctx.separated(&self.options, ", "); - let between = if self.arguments.is_empty() || self.options.is_empty() { - "" - } else { - ", " - }; - + textify_function_call_prefix( + self.function_reference, + &self.arguments, + &self.options, + ctx, + w, + )?; let output = OutputType(self.output_type.as_ref()); let output_type = ctx.display(&output); - write!(w, "{name_and_anchor}({args}{between}{options}) over(")?; + write!(w, " over(")?; let mut named_args: Vec = Vec::new(); + // phase= is always written, even when Unspecified: unlike invocation=, + // the parser requires a phase= argument to be present in over(...). let mut phase_str = String::new(); - match aggregation_phase_str(self.phase) { - Ok(s) => textify_enum(s, ctx, &mut phase_str)?, - Err(e) => write!(phase_str, "{}", ctx.failure(e))?, - } + textify_i32_enum::(self.phase, "phase", ctx, &mut phase_str)?; named_args.push(format!("phase={phase_str}")); // order= is omitted when there are no sort fields. A single sort field @@ -671,10 +642,12 @@ impl Textify for WindowFunction { if self.invocation != AggregationInvocation::Unspecified as i32 { let mut invocation_str = String::new(); - match aggregation_invocation_str(self.invocation) { - Ok(s) => textify_enum(s, ctx, &mut invocation_str)?, - Err(e) => write!(invocation_str, "{}", ctx.failure(e))?, - } + textify_i32_enum::( + self.invocation, + "invocation", + ctx, + &mut invocation_str, + )?; named_args.push(format!("invocation={invocation_str}")); } @@ -688,7 +661,25 @@ impl Textify for WindowFunction { if !matches!(bounds_type, Ok(window_function::BoundsType::Unspecified)) || has_bounds { let keyword = match bounds_type { Ok(window_function::BoundsType::Rows) => "rows", - Ok(window_function::BoundsType::Range) => "range", + Ok(window_function::BoundsType::Range) => { + // The parser rejects range= frames unless there is + // exactly one order= field; enforce the same invariant + // here so textified output always re-parses. + if self.sorts.len() != 1 { + ctx.push_error( + PlanError::invalid( + "WindowFunction", + Some("bounds_type"), + format!( + "range frame requires exactly one order= field, found {}", + self.sorts.len() + ), + ) + .into(), + ); + } + "range" + } Ok(window_function::BoundsType::Unspecified) => { // bounds_type is required whenever a frame is present. ctx.push_error( @@ -714,9 +705,9 @@ impl Textify for WindowFunction { } }; let mut lower_str = String::new(); - textify_window_bound(self.lower_bound.as_ref(), ctx, &mut lower_str)?; + textify_window_bound(self.lower_bound.as_ref(), &mut lower_str)?; let mut upper_str = String::new(); - textify_window_bound(self.upper_bound.as_ref(), ctx, &mut upper_str)?; + textify_window_bound(self.upper_bound.as_ref(), &mut upper_str)?; named_args.push(format!("{keyword}=({lower_str}, {upper_str})")); } @@ -840,26 +831,16 @@ impl Textify for AggregateFunction { } fn textify(&self, ctx: &S, w: &mut W) -> fmt::Result { - // Similar to ScalarFunction textification - let name_and_anchor = - NamedAnchor::lookup(ctx, ExtensionKind::Function, self.function_reference); - let name_and_anchor = ctx.display(&name_and_anchor); - - let args = ctx.separated(&self.arguments, ", "); - let options = ctx.separated(&self.options, ", "); - let between = if self.arguments.is_empty() || self.options.is_empty() { - "" - } else { - ", " - }; + textify_function_call_prefix( + self.function_reference, + &self.arguments, + &self.options, + ctx, + w, + )?; let output = OutputType(self.output_type.as_ref()); - let output_type = ctx.display(&output); - - write!( - w, - "{name_and_anchor}({args}{between}{options}){output_type}" - ) + write!(w, "{}", ctx.display(&output)) } } @@ -867,6 +848,7 @@ impl Textify for AggregateFunction { mod tests { use substrait::proto::Type; use substrait::proto::expression::{cast, if_then}; + use substrait::proto::sort_field::{SortDirection, SortKind}; use substrait::proto::r#type::{Boolean, I16, I32, I64, Kind, Nullability, UserDefined}; use super::*; diff --git a/src/textify/rels.rs b/src/textify/rels.rs index 433cf037..0559ad08 100644 --- a/src/textify/rels.rs +++ b/src/textify/rels.rs @@ -5,6 +5,7 @@ use std::fmt; use std::fmt::Debug; use prost::{Message, UnknownEnumValue}; +use substrait::proto::aggregate_function::AggregationInvocation; use substrait::proto::fetch_rel::CountMode; use substrait::proto::plan_rel::RelType as PlanRelType; use substrait::proto::read_rel::ReadType; @@ -12,9 +13,10 @@ use substrait::proto::rel::RelType; use substrait::proto::rel_common::EmitKind; use substrait::proto::sort_field::{SortDirection, SortKind}; use substrait::proto::{ - AggregateFunction, AggregateRel, CrossRel, Expression, ExtensionLeafRel, ExtensionMultiRel, - ExtensionSingleRel, FetchRel, FilterRel, JoinRel, NamedStruct, PlanRel, ProjectRel, ReadRel, - Rel, RelCommon, RelRoot, SetRel, SortField, SortRel, Type, join_rel, set_rel, + AggregateFunction, AggregateRel, AggregationPhase, CrossRel, Expression, ExtensionLeafRel, + ExtensionMultiRel, ExtensionSingleRel, FetchRel, FilterRel, JoinRel, NamedStruct, PlanRel, + ProjectRel, ReadRel, Rel, RelCommon, RelRoot, SetRel, SortField, SortRel, Type, join_rel, + set_rel, }; use super::addenda::AddendumLines; @@ -1229,6 +1231,30 @@ impl ValueEnum for set_rel::SetOp { } } +impl ValueEnum for AggregationPhase { + fn as_enum_str(&self) -> Result, PlanError> { + let s = match self { + AggregationPhase::Unspecified => "Unspecified", + AggregationPhase::InitialToIntermediate => "InitialToIntermediate", + AggregationPhase::IntermediateToIntermediate => "IntermediateToIntermediate", + AggregationPhase::InitialToResult => "InitialToResult", + AggregationPhase::IntermediateToResult => "IntermediateToResult", + }; + Ok(Cow::Borrowed(s)) + } +} + +impl ValueEnum for AggregationInvocation { + fn as_enum_str(&self) -> Result, PlanError> { + let s = match self { + AggregationInvocation::Unspecified => "Unspecified", + AggregationInvocation::All => "All", + AggregationInvocation::Distinct => "Distinct", + }; + Ok(Cow::Borrowed(s)) + } +} + impl<'a> Textify for NamedArg<'a> { fn name() -> &'static str { "NamedArg" diff --git a/tests/plan_roundtrip.rs b/tests/plan_roundtrip.rs index 6bd56711..65002d3b 100644 --- a/tests/plan_roundtrip.rs +++ b/tests/plan_roundtrip.rs @@ -1219,3 +1219,39 @@ Root[s] assert!(Parser::parse(plan).is_err()); } + +/// A `range=` frame (with exactly one `order=` field), a positive bound +/// offset (`FOLLOWING`), and a non-default `invocation=&All` all round-trip. +#[test] +fn test_window_function_range_following_and_invocation_all_roundtrip() { + let plan = r#"=== Extensions +URNs: + @ 1: https://github.com/substrait-io/substrait/blob/main/extensions/functions_aggregate.yaml +Functions: + # 10 @ 1: sum + +=== Plan +Root[s] + Project[sum($0) over(phase=&InitialToResult, order=($0,&AscNullsLast), invocation=&All, range=(_, 5)):fp64?] + Read[t => a:fp64]"#; + + roundtrip_plan(plan); +} + +/// `phase=` is a required named argument in `over(...)`; a window function +/// call missing it fails to parse at the full-plan level. +#[test] +fn test_window_function_missing_phase_fails_at_plan_level() { + let plan = r#"=== Extensions +URNs: + @ 1: https://github.com/substrait-io/substrait/blob/main/extensions/functions_arithmetic.yaml +Functions: + # 10 @ 1: row_number + +=== Plan +Root[r] + Project[row_number() over(partition=($0)):i64] + Read[t => a:i32]"#; + + assert!(Parser::parse(plan).is_err()); +} From e0829f21d92a731521a3d66b6cc30a3ea344448f Mon Sep 17 00:00:00 2001 From: "gordon.hamilton" Date: Thu, 16 Jul 2026 09:21:35 -0400 Subject: [PATCH 3/7] feat: fixing overflow bug --- GRAMMAR.md | 17 +- src/parser/expression_grammar.pest | 4 +- src/parser/expressions.rs | 152 ++++++++--------- src/textify/expressions.rs | 257 +++++++++++++++++++---------- tests/plan_roundtrip.rs | 46 +++++- 5 files changed, 300 insertions(+), 176 deletions(-) diff --git a/GRAMMAR.md b/GRAMMAR.md index 5c147311..bc056b38 100644 --- a/GRAMMAR.md +++ b/GRAMMAR.md @@ -509,10 +509,19 @@ A window function computes a value over a "window" of rows related to the curren `window_function := function_signature anchor? urn_anchor? "(" (expression ("," expression)*)? ")" "over(" (window_named_arg ("," window_named_arg)*)? ")" ":" type` -`window_named_arg` is one of: +```text +window_named_arg := "partition=" "(" expression ("," expression)* ")" + / "order=" sort_field + / "order=" "(" sort_field ("," sort_field)+ ")" + / "invocation=" enum + / "phase=" enum + / ("rows=" / "range=") "(" window_bound "," window_bound ")" +sort_field := "(" reference "," enum ")" +window_bound := integer / "_" +``` - `partition=(expression, ...)` - partitioning expressions -- `order=sort_field` or `order=(sort_field, sort_field, ...)` - ordering field(s); a single sort field is written bare (e.g. `order=($1,&AscNullsLast)`), two or more are wrapped in a parenthesized list of tuples (e.g. `order=(($1,&AscNullsLast),($2,&DescNullsLast))`) +- `order=sort_field` or `order=(sort_field, sort_field, ...)` - ordering field(s); a single sort field is written bare (e.g. `order=($1, &AscNullsLast)`), two or more are wrapped in a parenthesized list of tuples (e.g. `order=(($1, &AscNullsLast), ($2, &DescNullsLast))`) - `invocation=&Distinct` / `invocation=&All` - aggregation invocation - `phase=&InitialToResult` (etc.) - aggregation phase - `rows=(lower, upper)` / `range=(lower, upper)` - the window frame, each bound is an integer (negative = preceding, @@ -534,7 +543,7 @@ Functions: === Plan Root[a, b, s] - Project[$0, $1, sum($1) over(phase=&InitialToResult, order=($1,&AscNullsLast), invocation=&Distinct, partition=($0), rows=(-3, 0)):fp64?] + Project[$0, $1, sum($1) over(phase=&InitialToResult, order=($1, &AscNullsLast), invocation=&Distinct, partition=($0), rows=(-3, 0)):fp64?] Read[t => a:i32, b:fp64] # "#; # @@ -556,7 +565,7 @@ Functions: === Plan Root[a, r] - Project[$0, row_number() over(phase=&InitialToResult, order=($0,&AscNullsLast), partition=($0)):i64] + Project[$0, row_number() over(phase=&InitialToResult, order=($0, &AscNullsLast), partition=($0)):i64] Read[t => a:i32] # "#; # diff --git a/src/parser/expression_grammar.pest b/src/parser/expression_grammar.pest index 5562ff81..25b199c1 100644 --- a/src/parser/expression_grammar.pest +++ b/src/parser/expression_grammar.pest @@ -199,8 +199,8 @@ cast_expression = { "(" ~ sp ~ expression ~ sp ~ ")" ~ sp ~ "::" ~ sp ~ cast_fai // - `phase=&...` - required aggregation phase // - `invocation=&...` - optional invocation (e.g. `&All`, `&Distinct`) window_partition_arg = { "partition" ~ sp ~ "=" ~ sp ~ "(" ~ sp ~ expression_list? ~ sp ~ ")" } -order_value = { sort_field | ("(" ~ sp ~ sort_field ~ (sp ~ "," ~ sp ~ sort_field)+ ~ sp ~ ")") } -window_order_arg = { "order" ~ sp ~ "=" ~ sp ~ order_value } +window_sort_order = { sort_field | ("(" ~ sp ~ sort_field ~ (sp ~ "," ~ sp ~ sort_field)+ ~ sp ~ ")") } +window_order_arg = { "order" ~ sp ~ "=" ~ sp ~ window_sort_order } window_bound = { integer | empty } window_frame_kind = { "rows" | "range" } window_frame_arg = { window_frame_kind ~ sp ~ "=" ~ sp ~ "(" ~ sp ~ window_bound ~ sp ~ "," ~ sp ~ window_bound ~ sp ~ ")" } diff --git a/src/parser/expressions.rs b/src/parser/expressions.rs index 23e15cf7..28c6aab0 100644 --- a/src/parser/expressions.rs +++ b/src/parser/expressions.rs @@ -1,3 +1,5 @@ +use std::fmt; + use chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime}; use substrait::proto::aggregate_function::AggregationInvocation; use substrait::proto::aggregate_rel::Measure; @@ -401,14 +403,14 @@ impl ScopedParsePair for Literal { } } -/// The shared prefix of a scalar/window function call. +/// The shared prefix of function calls. struct FunctionHead { name: CompoundName, anchor: Option, arguments: Vec, } -/// Parse the `CompoundName ~ anchor? ~ urn_anchor? ~ argument_list` prefix shared by `ScalarFunction` and `WindowFunction`. +/// Parse the `CompoundName ~ anchor? ~ urn_anchor? ~ argument_list` prefix of functions. fn parse_function_head( extensions: &SimpleExtensions, iter: &mut RuleIter<'_>, @@ -492,71 +494,41 @@ impl ScopedParsePair for ScalarFunction { /// `_` (empty) means unbounded, represented explicitly as `Some(Bound { kind: /// Some(Unbounded{}) })` rather than `None`. `0` maps to `CurrentRow`; /// negative/positive integers map to `Preceding`/`Following`. -fn parse_window_bound(pair: pest::iterators::Pair) -> Option { +fn parse_window_bound( + pair: pest::iterators::Pair, +) -> Result, MessageParseError> { assert_eq!(pair.as_rule(), Rule::window_bound); let inner = unwrap_single_pair(pair); match inner.as_rule() { - Rule::empty => Some(window_function::Bound { + Rule::empty => Ok(Some(window_function::Bound { kind: Some(bound::Kind::Unbounded(bound::Unbounded {})), - }), + })), Rule::integer => { - let offset: i64 = inner.as_str().parse().unwrap(); + let out_of_range = |value: &dyn fmt::Display| { + MessageParseError::invalid( + "window_bound", + inner.as_span(), + format!("Window bound offset '{value}' is out of range"), + ) + }; + let offset: i64 = inner + .as_str() + .parse() + .map_err(|_| out_of_range(&inner.as_str()))?; let kind = match offset { 0 => bound::Kind::CurrentRow(bound::CurrentRow {}), n if n > 0 => bound::Kind::Following(bound::Following { offset: n }), - n => bound::Kind::Preceding(bound::Preceding { offset: -n }), + n => { + let offset = n.checked_neg().ok_or_else(|| out_of_range(&n))?; + bound::Kind::Preceding(bound::Preceding { offset }) + } }; - Some(window_function::Bound { kind: Some(kind) }) + Ok(Some(window_function::Bound { kind: Some(kind) })) } other => unreachable!("Grammar guarantees window_bound is integer or empty, got {other:?}"), } } -/// Parse an `enum_value` pair (`&Identifier`) into the `i32` discriminant of -/// one of `variants` (name, discriminant pairs), for enums that are looked up by name. -fn parse_enum_value_by_name( - pair: pest::iterators::Pair, - type_name: &'static str, - variants: &[(&str, i32)], -) -> Result { - assert_eq!(pair.as_rule(), Rule::enum_value); - let name = pair.as_str().trim_start_matches('&'); - variants - .iter() - .find(|(variant, _)| *variant == name) - .map(|(_, value)| *value) - .ok_or_else(|| { - MessageParseError::invalid( - type_name, - pair.as_span(), - format!("Unknown {type_name}: {name}"), - ) - }) -} - -const AGGREGATION_PHASE_VARIANTS: &[(&str, i32)] = &[ - ("Unspecified", AggregationPhase::Unspecified as i32), - ( - "InitialToIntermediate", - AggregationPhase::InitialToIntermediate as i32, - ), - ( - "IntermediateToIntermediate", - AggregationPhase::IntermediateToIntermediate as i32, - ), - ("InitialToResult", AggregationPhase::InitialToResult as i32), - ( - "IntermediateToResult", - AggregationPhase::IntermediateToResult as i32, - ), -]; - -const AGGREGATION_INVOCATION_VARIANTS: &[(&str, i32)] = &[ - ("Unspecified", AggregationInvocation::Unspecified as i32), - ("All", AggregationInvocation::All as i32), - ("Distinct", AggregationInvocation::Distinct as i32), -]; - impl ScopedParsePair for WindowFunction { fn rule() -> Rule { Rule::window_function_call @@ -608,9 +580,9 @@ impl ScopedParsePair for WindowFunction { } Rule::window_order_arg => { let mut order_iter = RuleIter::from(inner.into_inner()); - let order_value = order_iter.pop(Rule::order_value); + let sort_order = order_iter.pop(Rule::window_sort_order); order_iter.done(); - for sf in order_value.into_inner() { + for sf in sort_order.into_inner() { sorts.push(SortField::parse_pair(extensions, sf)?); } } @@ -627,28 +599,46 @@ impl ScopedParsePair for WindowFunction { "Grammar guarantees window_frame_kind is rows or range, got {other:?}" ), }; - lower_bound = parse_window_bound(lower); - upper_bound = parse_window_bound(upper); + lower_bound = parse_window_bound(lower)?; + upper_bound = parse_window_bound(upper)?; } Rule::window_phase_arg => { let mut phase_iter = RuleIter::from(inner.into_inner()); let enum_pair = phase_iter.pop(Rule::enum_value); phase_iter.done(); - phase = Some(parse_enum_value_by_name( - enum_pair, - "AggregationPhase", - AGGREGATION_PHASE_VARIANTS, - )?); + phase = Some(match enum_pair.as_str().trim_start_matches('&') { + "Unspecified" => AggregationPhase::Unspecified as i32, + "InitialToIntermediate" => AggregationPhase::InitialToIntermediate as i32, + "IntermediateToIntermediate" => { + AggregationPhase::IntermediateToIntermediate as i32 + } + "InitialToResult" => AggregationPhase::InitialToResult as i32, + "IntermediateToResult" => AggregationPhase::IntermediateToResult as i32, + other => { + return Err(MessageParseError::invalid( + "AggregationPhase", + enum_pair.as_span(), + format!("Unknown AggregationPhase: {other}"), + )); + } + }); } Rule::window_invocation_arg => { let mut inv_iter = RuleIter::from(inner.into_inner()); let enum_pair = inv_iter.pop(Rule::enum_value); inv_iter.done(); - invocation = parse_enum_value_by_name( - enum_pair, - "AggregationInvocation", - AGGREGATION_INVOCATION_VARIANTS, - )?; + invocation = match enum_pair.as_str().trim_start_matches('&') { + "Unspecified" => AggregationInvocation::Unspecified as i32, + "All" => AggregationInvocation::All as i32, + "Distinct" => AggregationInvocation::Distinct as i32, + other => { + return Err(MessageParseError::invalid( + "AggregationInvocation", + enum_pair.as_span(), + format!("Unknown AggregationInvocation: {other}"), + )); + } + }; } other => { unreachable!("Grammar guarantees window_named_arg alternatives, got {other:?}") @@ -1403,7 +1393,7 @@ mod tests { assert_eq!(pairs.as_str(), "equal:any_any"); } - fn make_extensions() -> SimpleExtensions { + fn make_extensions_for_fn_tests() -> SimpleExtensions { let mut exts = SimpleExtensions::default(); exts.add_extension_urn("urn".to_string(), 1).unwrap(); exts.add_extension( @@ -1433,7 +1423,7 @@ mod tests { #[test] fn test_scalar_function_full_compound_name() { // Full compound name without anchor - let exts = make_extensions(); + let exts = make_extensions_for_fn_tests(); let pair = parse_exact(Rule::function_call, "equal:any_any($0, $1):boolean"); let f = ScalarFunction::parse_pair(&exts, pair).unwrap(); assert_eq!(f.function_reference, 1); @@ -1446,7 +1436,7 @@ mod tests { #[test] fn test_scalar_function_second_overload() { - let exts = make_extensions(); + let exts = make_extensions_for_fn_tests(); let pair = parse_exact(Rule::function_call, "equal:str_str($0, $1):boolean"); let f = ScalarFunction::parse_pair(&exts, pair).unwrap(); @@ -1457,7 +1447,7 @@ mod tests { #[test] fn test_scalar_function_base_name_unique_overload() { // "add" has only one overload; base-name lookup should succeed - let exts = make_extensions(); + let exts = make_extensions_for_fn_tests(); let pair = parse_exact(Rule::function_call, "add($0, $1):i64"); let f = ScalarFunction::parse_pair(&exts, pair).unwrap(); @@ -1472,7 +1462,7 @@ mod tests { #[test] fn test_scalar_function_base_name_ambiguous_fails() { // "equal" has two overloads; base-name lookup should fail - let exts = make_extensions(); + let exts = make_extensions_for_fn_tests(); let pair = parse_exact(Rule::function_call, "equal($0, $1):boolean"); let result = ScalarFunction::parse_pair(&exts, pair); assert!(result.is_err(), "ambiguous base name should fail"); @@ -1480,7 +1470,7 @@ mod tests { #[test] fn test_scalar_function_compound_name_with_anchor() { - let exts = make_extensions(); + let exts = make_extensions_for_fn_tests(); let pair = parse_exact(Rule::function_call, "equal:any_any#1($0, $1):boolean"); let f = ScalarFunction::parse_pair(&exts, pair).unwrap(); assert_eq!(f.function_reference, 1); @@ -1490,7 +1480,7 @@ mod tests { #[test] fn test_scalar_function_base_name_with_anchor() { // Base name + explicit anchor should resolve (anchor 1 stores equal:any_any) - let exts = make_extensions(); + let exts = make_extensions_for_fn_tests(); let pair = parse_exact(Rule::function_call, "equal#1($0, $1):boolean"); let f = ScalarFunction::parse_pair(&exts, pair).unwrap(); assert_eq!(f.function_reference, 1); @@ -1499,7 +1489,7 @@ mod tests { #[test] fn test_scalar_function_wrong_name_for_anchor_fails() { - let exts = make_extensions(); + let exts = make_extensions_for_fn_tests(); let pair = parse_exact(Rule::function_call, "like#1($0):boolean"); let result = ScalarFunction::parse_pair(&exts, pair); assert!(result.is_err(), "mismatched name/anchor should fail"); @@ -1698,7 +1688,7 @@ mod tests { #[test] fn test_window_function_full() { - let exts = make_extensions(); + let exts = make_extensions_for_fn_tests(); let pair = parse_exact( Rule::window_function_call, "add:i64_i64($0, $1) over(partition=($0), order=($1,&AscNullsLast), invocation=&Distinct, rows=(-3, 0), phase=&InitialToResult):i64", @@ -1727,7 +1717,7 @@ mod tests { #[test] fn test_window_function_unbounded() { - let exts = make_extensions(); + let exts = make_extensions_for_fn_tests(); let pair = parse_exact( Rule::window_function_call, "add:i64_i64($0, $1) over(phase=&InitialToResult):i64", @@ -1746,7 +1736,7 @@ mod tests { #[test] fn test_window_function_unbounded_lower_bound() { - let exts = make_extensions(); + let exts = make_extensions_for_fn_tests(); let pair = parse_exact( Rule::window_function_call, "add:i64_i64($0, $1) over(rows=(_, 5), phase=&InitialToResult):i64", @@ -1768,7 +1758,7 @@ mod tests { #[test] fn test_window_function_missing_phase_fails() { - let exts = make_extensions(); + let exts = make_extensions_for_fn_tests(); let pair = parse_exact( Rule::window_function_call, "add:i64_i64($0, $1) over(partition=($0)):i64", @@ -1779,7 +1769,7 @@ mod tests { #[test] fn test_window_function_range_with_multiple_order_fields_fails() { - let exts = make_extensions(); + let exts = make_extensions_for_fn_tests(); let pair = parse_exact( Rule::window_function_call, "add:i64_i64($0, $1) over(order=(($0,&AscNullsLast),($1,&AscNullsLast)), range=(_, 0), phase=&InitialToResult):i64", @@ -1793,7 +1783,7 @@ mod tests { #[test] fn test_window_function_call_in_expression() { - let exts = make_extensions(); + let exts = make_extensions_for_fn_tests(); let pair = parse_exact( Rule::expression, "add:i64_i64($0, $1) over(phase=&InitialToResult):i64", diff --git a/src/textify/expressions.rs b/src/textify/expressions.rs index a7491f17..f447ba5c 100644 --- a/src/textify/expressions.rs +++ b/src/textify/expressions.rs @@ -1,4 +1,5 @@ -use std::fmt::{self, Write as _}; +use std::borrow::Cow; +use std::fmt; use chrono::{DateTime, NaiveDate}; use expr::RexType; @@ -18,7 +19,7 @@ use substrait::proto::{ use super::{PlanError, Scope, Textify, Visibility}; use crate::extensions::simple::ExtensionKind; -use crate::textify::rels::ValueEnum; +use crate::textify::rels::{Arguments, NamedArg, Value, ValueEnum}; use crate::textify::types::{Name, NamedAnchor, OutputType, escaped}; // …(…) for function call @@ -540,57 +541,46 @@ impl Textify for IfThen { } } -/// Write a single sort field as `($ref,&Direction)`, for use in `order=`. -fn textify_sort_field(sf: &SortField, ctx: &S, w: &mut W) -> fmt::Result { - let expr = ctx.expect(sf.expr.as_ref()); - write!(w, "({expr},")?; - let result = match sf.sort_kind.as_ref() { - Some(sk) => sk.as_enum_str(), - None => Err(PlanError::invalid( - "SortField", - Some("sort_kind"), - "Missing sort_kind", +fn window_enum_value<'a, T: TryFrom + ValueEnum>( + raw: i32, + field_name: &'static str, +) -> Value<'a> { + match T::try_from(raw) { + Ok(v) => match v.as_enum_str() { + Ok(s) => Value::Enum(s), + Err(e) => Value::Missing(e), + }, + Err(_) => Value::Missing(PlanError::invalid( + "WindowFunction", + Some(field_name), + format!("Unknown {field_name}: {raw}"), )), - }; - match result { - Ok(s) => textify_enum(&s, ctx, w)?, - Err(e) => write!(w, "{}", ctx.failure(e))?, } - write!(w, ")") } -/// Resolve a raw `i32` discriminant to `T`, then render its name via `ValueEnum`, for enum fields. -fn textify_i32_enum(raw: i32, type_name: &'static str, ctx: &S, w: &mut W) -> fmt::Result -where - T: TryFrom + ValueEnum, - S: Scope, - W: fmt::Write, -{ - let result = match T::try_from(raw) { - Ok(v) => v.as_enum_str(), - Err(_) => Err(PlanError::invalid( - "WindowFunction", - Some(type_name), - format!("Unknown {type_name}: {raw}"), - )), - }; - match result { - Ok(s) => textify_enum(&s, ctx, w), - Err(e) => write!(w, "{}", ctx.failure(e)), +fn window_sort_order<'a>(sorts: &'a [SortField]) -> Value<'a> { + match sorts { + [single] => Value::from(single), + many => Value::Tuple(many.iter().map(Value::from).collect()), } } -/// Write a window bound as an integer offset (negative preceding, positive -/// following, `0` for the current row) or `_` for unbounded/unspecified. -fn textify_window_bound( +fn window_bound_value<'a>( bound: Option<&window_function::Bound>, - w: &mut W, -) -> fmt::Result { + field_name: &'static str, +) -> Value<'a> { match bound.and_then(|b| b.kind.as_ref()) { - None | Some(bound::Kind::Unbounded(_)) => write!(w, "_"), - Some(bound::Kind::CurrentRow(_)) => write!(w, "0"), - Some(bound::Kind::Preceding(p)) => write!(w, "{}", -p.offset), - Some(bound::Kind::Following(f)) => write!(w, "{}", f.offset), + None | Some(bound::Kind::Unbounded(_)) => Value::EmptyGroup, + Some(bound::Kind::CurrentRow(_)) => Value::Integer(0), + Some(bound::Kind::Preceding(p)) => match p.offset.checked_neg() { + Some(offset) => Value::Integer(offset), + None => Value::Missing(PlanError::invalid( + "WindowFunction", + Some(field_name), + format!("Window bound offset {} cannot be negated", p.offset), + )), + }, + Some(bound::Kind::Following(f)) => Value::Integer(f.offset), } } @@ -610,50 +600,35 @@ impl Textify for WindowFunction { let output = OutputType(self.output_type.as_ref()); let output_type = ctx.display(&output); - write!(w, " over(")?; - - let mut named_args: Vec = Vec::new(); + let mut named_args: Vec = Vec::new(); // phase= is always written, even when Unspecified: unlike invocation=, // the parser requires a phase= argument to be present in over(...). - let mut phase_str = String::new(); - textify_i32_enum::(self.phase, "phase", ctx, &mut phase_str)?; - named_args.push(format!("phase={phase_str}")); + named_args.push(NamedArg { + name: Cow::Borrowed("phase"), + value: window_enum_value::(self.phase, "phase"), + }); - // order= is omitted when there are no sort fields. A single sort field - // is written bare; two or more are wrapped in a parenthesized list of - // tuples. + // order= is omitted when there are no sort fields. if !self.sorts.is_empty() { - let mut order_str = String::new(); - if self.sorts.len() == 1 { - textify_sort_field(&self.sorts[0], ctx, &mut order_str)?; - } else { - write!(order_str, "(")?; - for (i, sf) in self.sorts.iter().enumerate() { - if i > 0 { - write!(order_str, ",")?; - } - textify_sort_field(sf, ctx, &mut order_str)?; - } - write!(order_str, ")")?; - } - named_args.push(format!("order={order_str}")); + named_args.push(NamedArg { + name: Cow::Borrowed("order"), + value: window_sort_order(&self.sorts), + }); } if self.invocation != AggregationInvocation::Unspecified as i32 { - let mut invocation_str = String::new(); - textify_i32_enum::( - self.invocation, - "invocation", - ctx, - &mut invocation_str, - )?; - named_args.push(format!("invocation={invocation_str}")); + named_args.push(NamedArg { + name: Cow::Borrowed("invocation"), + value: window_enum_value::(self.invocation, "invocation"), + }); } if !self.partitions.is_empty() { - let parts = ctx.separated(&self.partitions, ", "); - named_args.push(format!("partition=({parts})")); + named_args.push(NamedArg { + name: Cow::Borrowed("partition"), + value: Value::Tuple(self.partitions.iter().map(Value::Expression).collect()), + }); } let bounds_type = window_function::BoundsType::try_from(self.bounds_type); @@ -704,14 +679,16 @@ impl Textify for WindowFunction { "rows" } }; - let mut lower_str = String::new(); - textify_window_bound(self.lower_bound.as_ref(), &mut lower_str)?; - let mut upper_str = String::new(); - textify_window_bound(self.upper_bound.as_ref(), &mut upper_str)?; - named_args.push(format!("{keyword}=({lower_str}, {upper_str})")); + let lower = window_bound_value(self.lower_bound.as_ref(), "lower_bound"); + let upper = window_bound_value(self.upper_bound.as_ref(), "upper_bound"); + named_args.push(NamedArg { + name: Cow::Borrowed(keyword), + value: Value::Tuple(vec![lower, upper]), + }); } - write!(w, "{}", named_args.join(", "))?; + write!(w, " over(")?; + Arguments::inline(vec![], named_args).textify(ctx, w)?; write!(w, "){output_type}") } } @@ -1374,7 +1351,7 @@ mod tests { }); assert_eq!( ctx.textify_no_errors(&f), - "sum() over(phase=&InitialToResult, order=($1,&AscNullsLast), invocation=&Distinct, partition=($0), rows=(-3, 0)):i16" + "sum() over(phase=&InitialToResult, order=($1, &AscNullsLast), invocation=&Distinct, partition=($0), rows=(-3, 0)):i16" ); } @@ -1390,7 +1367,7 @@ mod tests { ]; assert_eq!( ctx.textify_no_errors(&f), - "sum() over(phase=&InitialToResult, order=(($0,&AscNullsLast),($1,&DescNullsFirst))):i16" + "sum() over(phase=&InitialToResult, order=(($0, &AscNullsLast), ($1, &DescNullsFirst))):i16" ); } @@ -1427,4 +1404,114 @@ mod tests { assert_eq!(s, "sum() over(phase=&InitialToResult, rows=(0, _)):i16"); assert!(!errs.is_empty(), "expected a diagnostic about bounds_type"); } + + #[test] + fn test_window_function_i64_min_preceding_bound_does_not_panic() { + // i64::MIN has no positive counterpart, so negating it as a + // "preceding" offset must surface an accumulated error instead of + // panicking (or silently wrapping) via unary negation. + let ctx = TestContext::new() + .with_urn(1, "urn:example") + .with_function(1, 10, "sum"); + let mut f = base_window_function(); + f.bounds_type = window_function::BoundsType::Rows as i32; + f.lower_bound = Some(window_function::Bound { + kind: Some(bound::Kind::Preceding(bound::Preceding { + offset: i64::MIN, + })), + }); + f.upper_bound = Some(window_function::Bound { + kind: Some(bound::Kind::CurrentRow(bound::CurrentRow {})), + }); + let (s, errs) = ctx.textify(&f); + assert_eq!( + s, + "sum() over(phase=&InitialToResult, rows=(!{WindowFunction}, 0)):i16" + ); + assert!(!errs.is_empty(), "expected a diagnostic about the bound"); + } + + #[test] + fn test_window_function_i64_max_preceding_bound_roundtrips() { + // i64::MAX *does* have a negatable counterpart (-i64::MAX, which is + // one more than i64::MIN), so this boundary value must still + // textify cleanly rather than being caught by the i64::MIN check. + let ctx = TestContext::new() + .with_urn(1, "urn:example") + .with_function(1, 10, "sum"); + let mut f = base_window_function(); + f.bounds_type = window_function::BoundsType::Rows as i32; + f.lower_bound = Some(window_function::Bound { + kind: Some(bound::Kind::Preceding(bound::Preceding { + offset: i64::MAX, + })), + }); + f.upper_bound = Some(window_function::Bound { + kind: Some(bound::Kind::CurrentRow(bound::CurrentRow {})), + }); + assert_eq!( + ctx.textify_no_errors(&f), + "sum() over(phase=&InitialToResult, rows=(-9223372036854775807, 0)):i16" + ); + } + + #[test] + fn test_window_function_range_wrong_sort_count() { + // range= requires exactly one order= field; zero (or more than one) + // is still rendered best-effort but surfaces a diagnostic. + let ctx = TestContext::new() + .with_urn(1, "urn:example") + .with_function(1, 10, "sum"); + let mut f = base_window_function(); + f.bounds_type = window_function::BoundsType::Range as i32; + f.lower_bound = Some(window_function::Bound { + kind: Some(bound::Kind::Unbounded(bound::Unbounded {})), + }); + f.upper_bound = Some(window_function::Bound { + kind: Some(bound::Kind::CurrentRow(bound::CurrentRow {})), + }); + let (s, errs) = ctx.textify(&f); + assert_eq!(s, "sum() over(phase=&InitialToResult, range=(_, 0)):i16"); + assert!(!errs.is_empty(), "expected a diagnostic about sort count"); + } + + #[test] + fn test_window_function_unknown_bounds_type() { + // bounds_type 99 matches no BoundsType variant + let ctx = TestContext::new() + .with_urn(1, "urn:example") + .with_function(1, 10, "sum"); + let mut f = base_window_function(); + f.bounds_type = 99; + f.lower_bound = Some(window_function::Bound { + kind: Some(bound::Kind::CurrentRow(bound::CurrentRow {})), + }); + f.upper_bound = Some(window_function::Bound { + kind: Some(bound::Kind::CurrentRow(bound::CurrentRow {})), + }); + let (s, errs) = ctx.textify(&f); + assert_eq!(s, "sum() over(phase=&InitialToResult, rows=(0, 0)):i16"); + assert!(!errs.is_empty(), "expected a diagnostic about bounds_type"); + } + + #[test] + fn test_window_function_unknown_phase_and_invocation() { + // phase 99 and invocation 99 match no know n enum variant, so both + // render as the same generic `!{WindowFunction}` error token. + let ctx = TestContext::new() + .with_urn(1, "urn:example") + .with_function(1, 10, "sum"); + let mut f = base_window_function(); + f.phase = 99; + f.invocation = 99; + let (s, errs) = ctx.textify(&f); + assert_eq!( + s, + "sum() over(phase=!{WindowFunction}, invocation=!{WindowFunction}):i16" + ); + assert!( + !errs.is_empty(), + "expected diagnostics about phase/invocation" + ); + } } diff --git a/tests/plan_roundtrip.rs b/tests/plan_roundtrip.rs index 65002d3b..e5ab0dd7 100644 --- a/tests/plan_roundtrip.rs +++ b/tests/plan_roundtrip.rs @@ -1144,7 +1144,7 @@ Functions: === Plan Root[a, b, s] - Project[$0, $1, sum($1) over(phase=&InitialToResult, order=($1,&AscNullsLast), invocation=&Distinct, partition=($0), rows=(-3, 0)):fp64?] + Project[$0, $1, sum($1) over(phase=&InitialToResult, order=($1, &AscNullsLast), invocation=&Distinct, partition=($0), rows=(-3, 0)):fp64?] Read[t => a:i32, b:fp64]"#; roundtrip_plan(plan); @@ -1160,7 +1160,7 @@ Functions: === Plan Root[a, r] - Project[$0, row_number() over(phase=&InitialToResult, order=($0,&AscNullsLast), partition=($0)):i64] + Project[$0, row_number() over(phase=&InitialToResult, order=($0, &AscNullsLast), partition=($0)):i64] Read[t => a:i32]"#; roundtrip_plan(plan); @@ -1178,7 +1178,7 @@ Functions: === Plan Root[r] - Project[row_number() over(phase=&InitialToResult, order=($0,&AscNullsLast)):i64] + Project[row_number() over(phase=&InitialToResult, order=($0, &AscNullsLast)):i64] Read[t => a:i32]"#; roundtrip_plan(plan); @@ -1232,7 +1232,7 @@ Functions: === Plan Root[s] - Project[sum($0) over(phase=&InitialToResult, order=($0,&AscNullsLast), invocation=&All, range=(_, 5)):fp64?] + Project[sum($0) over(phase=&InitialToResult, order=($0, &AscNullsLast), invocation=&All, range=(_, 5)):fp64?] Read[t => a:fp64]"#; roundtrip_plan(plan); @@ -1255,3 +1255,41 @@ Root[r] assert!(Parser::parse(plan).is_err()); } + +/// A `rows=` lower bound of `i64::MIN` can't be negated into a valid +/// `Preceding.offset` (`i64`); this must fail to parse with an error rather +/// than panicking or overflowing. +#[test] +fn test_window_function_min_i64_bound_fails_to_parse() { + let plan = r#"=== Extensions +URNs: + @ 1: https://github.com/substrait-io/substrait/blob/main/extensions/functions_arithmetic.yaml +Functions: + # 10 @ 1: sum + +=== Plan +Root[r] + Project[sum($0) over(phase=&InitialToResult, rows=(-9223372036854775808, 0)):i64?] + Read[t => a:i32]"#; + + assert!(Parser::parse(plan).is_err()); +} + +/// A `rows=` lower bound of `-i64::MAX` (one more than `i64::MIN`) *can* be +/// negated into a valid `Preceding.offset`, so it must parse and round-trip +/// successfully, unlike `i64::MIN` above. +#[test] +fn test_window_function_max_i64_bound_roundtrip() { + let plan = r#"=== Extensions +URNs: + @ 1: https://github.com/substrait-io/substrait/blob/main/extensions/functions_arithmetic.yaml +Functions: + # 10 @ 1: sum + +=== Plan +Root[r] + Project[sum($0) over(phase=&InitialToResult, rows=(-9223372036854775807, 0)):i64?] + Read[t => a:i32]"#; + + roundtrip_plan(plan); +} From 63266404a85972049c1a8d5d54519ac91072cf2a Mon Sep 17 00:00:00 2001 From: "gordon.hamilton" Date: Thu, 16 Jul 2026 12:56:04 -0400 Subject: [PATCH 4/7] feat: additional fixes --- GRAMMAR.md | 3 +- src/parser/expressions.rs | 38 ++++++++++++++++------ src/textify/expressions.rs | 66 ++++++++++++++++++++++++++------------ src/textify/rels.rs | 15 ++++++--- tests/plan_roundtrip.rs | 53 ++++++++++++++++++++++++++++++ 5 files changed, 140 insertions(+), 35 deletions(-) diff --git a/GRAMMAR.md b/GRAMMAR.md index bc056b38..4b5be85f 100644 --- a/GRAMMAR.md +++ b/GRAMMAR.md @@ -516,10 +516,11 @@ window_named_arg := "partition=" "(" expression ("," expression)* ")" / "invocation=" enum / "phase=" enum / ("rows=" / "range=") "(" window_bound "," window_bound ")" -sort_field := "(" reference "," enum ")" window_bound := integer / "_" ``` +`sort_field` is the same production used by the Sort relation (see below): `"(" reference "," sort_direction ")"`. + - `partition=(expression, ...)` - partitioning expressions - `order=sort_field` or `order=(sort_field, sort_field, ...)` - ordering field(s); a single sort field is written bare (e.g. `order=($1, &AscNullsLast)`), two or more are wrapped in a parenthesized list of tuples (e.g. `order=(($1, &AscNullsLast), ($2, &DescNullsLast))`) - `invocation=&Distinct` / `invocation=&All` - aggregation invocation diff --git a/src/parser/expressions.rs b/src/parser/expressions.rs index 28c6aab0..c1f10d01 100644 --- a/src/parser/expressions.rs +++ b/src/parser/expressions.rs @@ -1,5 +1,3 @@ -use std::fmt; - use chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime}; use substrait::proto::aggregate_function::AggregationInvocation; use substrait::proto::aggregate_rel::Measure; @@ -504,22 +502,24 @@ fn parse_window_bound( kind: Some(bound::Kind::Unbounded(bound::Unbounded {})), })), Rule::integer => { - let out_of_range = |value: &dyn fmt::Display| { + let offset: i64 = inner.as_str().parse().map_err(|_| { MessageParseError::invalid( "window_bound", inner.as_span(), - format!("Window bound offset '{value}' is out of range"), + format!("Window bound offset '{}' is out of range", inner.as_str()), ) - }; - let offset: i64 = inner - .as_str() - .parse() - .map_err(|_| out_of_range(&inner.as_str()))?; + })?; let kind = match offset { 0 => bound::Kind::CurrentRow(bound::CurrentRow {}), n if n > 0 => bound::Kind::Following(bound::Following { offset: n }), n => { - let offset = n.checked_neg().ok_or_else(|| out_of_range(&n))?; + let offset = n.checked_neg().ok_or_else(|| { + MessageParseError::invalid( + "window_bound", + inner.as_span(), + format!("Window bound offset '{n}' is out of range"), + ) + })?; bound::Kind::Preceding(bound::Preceding { offset }) } }; @@ -1756,6 +1756,24 @@ mod tests { ); } + #[test] + fn test_window_function_bound_offset_too_many_digits_fails() { + // The grammar's `integer` rule permits arbitrarily many digits, so a + // literal wider than i64 range must be rejected by the `.parse()` + // call itself, not just by the separate i64::MIN-negation overflow + // path exercised elsewhere. + let exts = make_extensions_for_fn_tests(); + let pair = parse_exact( + Rule::window_function_call, + "add:i64_i64($0, $1) over(rows=(-99999999999999999999, 0), phase=&InitialToResult):i64", + ); + let result = WindowFunction::parse_pair(&exts, pair); + assert!( + result.is_err(), + "a bound literal wider than i64 range must be rejected" + ); + } + #[test] fn test_window_function_missing_phase_fails() { let exts = make_extensions_for_fn_tests(); diff --git a/src/textify/expressions.rs b/src/textify/expressions.rs index f447ba5c..b32549f7 100644 --- a/src/textify/expressions.rs +++ b/src/textify/expressions.rs @@ -19,7 +19,7 @@ use substrait::proto::{ use super::{PlanError, Scope, Textify, Visibility}; use crate::extensions::simple::ExtensionKind; -use crate::textify::rels::{Arguments, NamedArg, Value, ValueEnum}; +use crate::textify::rels::{Arguments, NamedArg, Value, ValueEnum, enum_str_value}; use crate::textify::types::{Name, NamedAnchor, OutputType, escaped}; // …(…) for function call @@ -546,10 +546,7 @@ fn window_enum_value<'a, T: TryFrom + ValueEnum>( field_name: &'static str, ) -> Value<'a> { match T::try_from(raw) { - Ok(v) => match v.as_enum_str() { - Ok(s) => Value::Enum(s), - Err(e) => Value::Missing(e), - }, + Ok(v) => enum_str_value(v.as_enum_str()), Err(_) => Value::Missing(PlanError::invalid( "WindowFunction", Some(field_name), @@ -572,14 +569,17 @@ fn window_bound_value<'a>( match bound.and_then(|b| b.kind.as_ref()) { None | Some(bound::Kind::Unbounded(_)) => Value::EmptyGroup, Some(bound::Kind::CurrentRow(_)) => Value::Integer(0), - Some(bound::Kind::Preceding(p)) => match p.offset.checked_neg() { - Some(offset) => Value::Integer(offset), - None => Value::Missing(PlanError::invalid( - "WindowFunction", - Some(field_name), - format!("Window bound offset {} cannot be negated", p.offset), - )), - }, + Some(bound::Kind::Preceding(p)) if p.offset < 0 => Value::Missing(PlanError::invalid( + "WindowFunction", + Some(field_name), + format!("Preceding bound offset {} must not be negative", p.offset), + )), + Some(bound::Kind::Preceding(p)) => Value::Integer(-p.offset), + Some(bound::Kind::Following(f)) if f.offset < 0 => Value::Missing(PlanError::invalid( + "WindowFunction", + Some(field_name), + format!("Following bound offset {} must not be negative", f.offset), + )), Some(bound::Kind::Following(f)) => Value::Integer(f.offset), } } @@ -1406,19 +1406,45 @@ mod tests { } #[test] - fn test_window_function_i64_min_preceding_bound_does_not_panic() { - // i64::MIN has no positive counterpart, so negating it as a - // "preceding" offset must surface an accumulated error instead of - // panicking (or silently wrapping) via unary negation. + fn test_window_function_negative_following_bound_surfaces_error() { + // A negative "following" offset would textify as a bare negative integer, which re-parses + // as a "preceding" bound of the opposite sign. Surface an accumulated + // error instead of silently flipping the bound's direction. let ctx = TestContext::new() .with_urn(1, "urn:example") .with_function(1, 10, "sum"); let mut f = base_window_function(); f.bounds_type = window_function::BoundsType::Rows as i32; f.lower_bound = Some(window_function::Bound { - kind: Some(bound::Kind::Preceding(bound::Preceding { - offset: i64::MIN, - })), + kind: Some(bound::Kind::CurrentRow(bound::CurrentRow {})), + }); + f.upper_bound = Some(window_function::Bound { + kind: Some(bound::Kind::Following(bound::Following { offset: -5 })), + }); + let (s, errs) = ctx.textify(&f); + assert_eq!( + s, + "sum() over(phase=&InitialToResult, rows=(0, !{WindowFunction})):i16" + ); + assert!(!errs.is_empty(), "expected a diagnostic about the bound"); + } + + #[test] + fn test_window_function_negative_preceding_bound_surfaces_error() { + // A negative "preceding" offset would be negated into a bare positive integer, which + // re-parses as a "following" bound of the opposite sign. Surface an + // accumulated error instead of silently flipping the bound's + // direction. The guard is magnitude-agnostic (`offset < 0`), so this + // also covers i64::MIN: negation is never attempted on a negative + // offset, so there's no separate overflow-at-the-boundary path to + // test beyond this. + let ctx = TestContext::new() + .with_urn(1, "urn:example") + .with_function(1, 10, "sum"); + let mut f = base_window_function(); + f.bounds_type = window_function::BoundsType::Rows as i32; + f.lower_bound = Some(window_function::Bound { + kind: Some(bound::Kind::Preceding(bound::Preceding { offset: -5 })), }); f.upper_bound = Some(window_function::Bound { kind: Some(bound::Kind::CurrentRow(bound::CurrentRow {})), diff --git a/src/textify/rels.rs b/src/textify/rels.rs index 0559ad08..07463575 100644 --- a/src/textify/rels.rs +++ b/src/textify/rels.rs @@ -1135,12 +1135,19 @@ impl<'a> From<&'a SortField> for Value<'a> { } } +/// Converts an [`ValueEnum::as_enum_str`] result into a [`Value`]. Shared by +/// the blanket `From<&T>` impl below and by callers that only have an owned +/// enum value (and so can't borrow it for the lifetime `Value<'a>` requires). +pub(crate) fn enum_str_value<'a>(result: Result, PlanError>) -> Value<'a> { + match result { + Ok(s) => Value::Enum(s), + Err(e) => Value::Missing(e), + } +} + impl<'a, T: ValueEnum + ?Sized> From<&'a T> for Value<'a> { fn from(enum_val: &'a T) -> Self { - match enum_val.as_enum_str() { - Ok(s) => Value::Enum(s), - Err(e) => Value::Missing(e), - } + enum_str_value(enum_val.as_enum_str()) } } diff --git a/tests/plan_roundtrip.rs b/tests/plan_roundtrip.rs index e5ab0dd7..99ba367e 100644 --- a/tests/plan_roundtrip.rs +++ b/tests/plan_roundtrip.rs @@ -1293,3 +1293,56 @@ Root[r] roundtrip_plan(plan); } + +/// `partition=(...)` with more than one expression round-trips, preserving +/// both order and count. +#[test] +fn test_window_function_multi_expression_partition_roundtrip() { + let plan = r#"=== Extensions +URNs: + @ 1: https://github.com/substrait-io/substrait/blob/main/extensions/functions_arithmetic.yaml +Functions: + # 10 @ 1: row_number + +=== Plan +Root[a, b, r] + Project[$0, $1, row_number() over(phase=&InitialToResult, partition=($0, $1)):i64] + Read[t => a:i32, b:i32]"#; + + roundtrip_plan(plan); +} + +/// A `rows=` frame that is unbounded on both sides round-trips. +#[test] +fn test_window_function_fully_unbounded_bounds_roundtrip() { + let plan = r#"=== Extensions +URNs: + @ 1: https://github.com/substrait-io/substrait/blob/main/extensions/functions_aggregate.yaml +Functions: + # 10 @ 1: sum + +=== Plan +Root[s] + Project[sum($0) over(phase=&InitialToResult, rows=(_, _)):i64?] + Read[t => a:i64]"#; + + roundtrip_plan(plan); +} + +/// A `rows=` frame with a bounded (`PRECEDING`) lower bound and an unbounded +/// upper bound round-trips, mirroring the lower-unbounded case above. +#[test] +fn test_window_function_unbounded_upper_bound_roundtrip() { + let plan = r#"=== Extensions +URNs: + @ 1: https://github.com/substrait-io/substrait/blob/main/extensions/functions_aggregate.yaml +Functions: + # 10 @ 1: sum + +=== Plan +Root[s] + Project[sum($0) over(phase=&InitialToResult, rows=(-3, _)):i64?] + Read[t => a:i64]"#; + + roundtrip_plan(plan); +} From 0faeea9bc1808953f545a526f3c5fcd619558e33 Mon Sep 17 00:00:00 2001 From: "gordon.hamilton" Date: Tue, 21 Jul 2026 12:34:45 -0400 Subject: [PATCH 5/7] feat: updating code based on comments --- GRAMMAR.md | 46 +- src/parser/expression_grammar.pest | 51 +- src/parser/expressions.rs | 758 +++++++++++++++++++++-------- src/textify/expressions.rs | 336 ++++++++----- src/textify/mod.rs | 1 + src/textify/rels.rs | 332 +------------ src/textify/values.rs | 339 +++++++++++++ tests/plan_roundtrip.rs | 6 +- 8 files changed, 1182 insertions(+), 687 deletions(-) create mode 100644 src/textify/values.rs diff --git a/GRAMMAR.md b/GRAMMAR.md index 4b5be85f..6b9226dd 100644 --- a/GRAMMAR.md +++ b/GRAMMAR.md @@ -507,28 +507,26 @@ A window function computes a value over a "window" of rows related to the curren #### Syntax -`window_function := function_signature anchor? urn_anchor? "(" (expression ("," expression)*)? ")" "over(" (window_named_arg ("," window_named_arg)*)? ")" ":" type` - -```text -window_named_arg := "partition=" "(" expression ("," expression)* ")" - / "order=" sort_field - / "order=" "(" sort_field ("," sort_field)+ ")" - / "invocation=" enum - / "phase=" enum - / ("rows=" / "range=") "(" window_bound "," window_bound ")" -window_bound := integer / "_" -``` - -`sort_field` is the same production used by the Sort relation (see below): `"(" reference "," sort_direction ")"`. - -- `partition=(expression, ...)` - partitioning expressions -- `order=sort_field` or `order=(sort_field, sort_field, ...)` - ordering field(s); a single sort field is written bare (e.g. `order=($1, &AscNullsLast)`), two or more are wrapped in a parenthesized list of tuples (e.g. `order=(($1, &AscNullsLast), ($2, &DescNullsLast))`) -- `invocation=&Distinct` / `invocation=&All` - aggregation invocation -- `phase=&InitialToResult` (etc.) - aggregation phase -- `rows=(lower, upper)` / `range=(lower, upper)` - the window frame, each bound is an integer (negative = preceding, - positive = following, `0` = current row) or `_` for unbounded - -`partition=`, `order=`, `invocation=`, and `rows=`/`range=` are all optional and are omitted when empty; `phase=` is always required and always printed, though this is enforced when parsing rather than by the grammar itself (the named-arg list as a whole is syntactically optional). A `range=` frame requires exactly one `order=` field. +`window_function := function_signature anchor? urn_anchor? "(" (expression ("," expression)*)? ")" "over(" window_named_arguments? ")" ":" type` + +Rather than a bespoke grammar production per named argument, `over(...)` reuses the same +generic [`named_arguments`](#arguments) production used elsewhere (`name "=" argument`, +comma-separated). Rust code (not the grammar) decides which names are allowed in `over(...)` +and what shape each value must have: + +- `partition=argument` - optional partitioning expression(s); a single field may be written + bare (`partition=$1`) or in a one-element tuple (`partition=($1,)`); multiple fields use a + multi-element tuple (`partition=($1, $2)`) +- `order=argument` - optional sort key(s); a single sort field is a bare 2-tuple of + `(reference, direction)` (e.g. `order=($3, &AscNullsLast)`); multiple sort fields are a + tuple of such 2-tuples (e.g. `order=(($3, &AscNullsLast), ($4, &DescNullsLast))`) +- `invocation=&Distinct` / `invocation=&All` - optional aggregation invocation +- `phase=&InitialToResult` (etc.) - required aggregation phase +- `rows=(lower, upper)` / `range=(lower, upper)` - optional, mutually exclusive window frame; + each bound is an integer (negative = preceding, positive = following, `0` = current row) or + `_` for unbounded + +`partition=`, `order=`, `invocation=`, and `rows=`/`range=` are all optional and are omitted when empty; `phase=` is always required and always printed, though this (along with the mutual exclusion of `rows=`/`range=`, and the requirement that a `range=` frame have exactly one `order=` field) is enforced when parsing rather than by the grammar itself (the named-arg list as a whole is syntactically optional). #### Examples @@ -544,7 +542,7 @@ Functions: === Plan Root[a, b, s] - Project[$0, $1, sum($1) over(phase=&InitialToResult, order=($1, &AscNullsLast), invocation=&Distinct, partition=($0), rows=(-3, 0)):fp64?] + Project[$0, $1, sum($1) over(phase=&InitialToResult, order=($1, &AscNullsLast), invocation=&Distinct, partition=$0, rows=(-3, 0)):fp64?] Read[t => a:i32, b:fp64] # "#; # @@ -566,7 +564,7 @@ Functions: === Plan Root[a, r] - Project[$0, row_number() over(phase=&InitialToResult, order=($0, &AscNullsLast), partition=($0)):i64] + Project[$0, row_number() over(phase=&InitialToResult, order=($0, &AscNullsLast), partition=$0):i64] Read[t => a:i32] # "#; # diff --git a/src/parser/expression_grammar.pest b/src/parser/expression_grammar.pest index 25b199c1..b074c698 100644 --- a/src/parser/expression_grammar.pest +++ b/src/parser/expression_grammar.pest @@ -146,19 +146,19 @@ type = { simple_type | compound_type | user_defined_type } argument_list = { "(" ~ (expression ~ (sp ~ "," ~ sp ~ expression)*)? ~ ")" } -// Scalar function call: +// An unresolved reference to a function: // - Compound function name (base name or base:signature, e.g. "equal" or "equal:any_any") // - Optional Anchor, e.g. #1 // - Optional URN Anchor, e.g. @1 +function_reference = { function_signature ~ sp ~ anchor? ~ sp ~ urn_anchor? } + +// Scalar function call: +// - A function_reference // - Arguments `()` are required, e.g. (1, 2, 3) // - Required output type annotation after closing paren, e.g. :i64 // (unambiguous because function_signature is atomic and ends before the opening paren) -function_call_prefix = _{ - function_signature ~ sp ~ anchor? ~ sp ~ urn_anchor? ~ sp ~ argument_list -} - function_call = { - function_call_prefix ~ ":" ~ sp ~ type + function_reference ~ sp ~ argument_list ~ ":" ~ sp ~ type } if_clause = { @@ -189,29 +189,32 @@ cast_expression = { "(" ~ sp ~ expression ~ sp ~ ")" ~ sp ~ "::" ~ sp ~ cast_fai // by a "over(...)" clause carrying window-specific named arguments. // Example: sum($0) over(partition=($1,$2), order=($3,&AscNullsLast), rows=(-3, 0), phase=&InitialToResult, invocation=&Distinct):fp64 // -// - `partition=(...)` - optional list of partitioning expressions -// - `order=...` - optional sort key(s); a single sort field is written bare -// (e.g. `order=($3,&AscNullsLast)`), two or more are wrapped in a list -// (e.g. `order=(($3,&AscNullsLast),($4,&DescNullsLast))`) -// - `rows=(lower,upper)` / `range=(lower,upper)` - optional frame bounds, encodes `bounds_type` with rows and range. -// Each bound is an integer offset (negative preceding, positive following, 0 for the current row) -// or `_` for unbounded/unspecified. +// - `partition=(...)` - optional list of partitioning expressions; a single +// partition field may be written bare (`partition=$1`) or in a one-element +// tuple (`partition=($1,)`); multiple fields use a multi-element tuple. +// - `order=...` - optional sort key(s); a single sort field is a bare 2-tuple +// of (reference, direction) (e.g. `order=($3,&AscNullsLast)`); multiple sort +// fields are a tuple of such 2-tuples (e.g. `order=(($3,&AscNullsLast),($4,&DescNullsLast))`) +// - `rows=(lower,upper)` / `range=(lower,upper)` - optional frame bounds (mutually +// exclusive), encodes `bounds_type` with rows and range. +// Each bound is an integer offset (negative preceding, positive following, 0 +// for the current row) or `_` for unbounded/unspecified. // - `phase=&...` - required aggregation phase // - `invocation=&...` - optional invocation (e.g. `&All`, `&Distinct`) -window_partition_arg = { "partition" ~ sp ~ "=" ~ sp ~ "(" ~ sp ~ expression_list? ~ sp ~ ")" } -window_sort_order = { sort_field | ("(" ~ sp ~ sort_field ~ (sp ~ "," ~ sp ~ sort_field)+ ~ sp ~ ")") } -window_order_arg = { "order" ~ sp ~ "=" ~ sp ~ window_sort_order } -window_bound = { integer | empty } -window_frame_kind = { "rows" | "range" } -window_frame_arg = { window_frame_kind ~ sp ~ "=" ~ sp ~ "(" ~ sp ~ window_bound ~ sp ~ "," ~ sp ~ window_bound ~ sp ~ ")" } -window_phase_arg = { "phase" ~ sp ~ "=" ~ sp ~ enum_value } -window_invocation_arg = { "invocation" ~ sp ~ "=" ~ sp ~ enum_value } - -window_named_arg = { window_partition_arg | window_order_arg | window_frame_arg | window_phase_arg | window_invocation_arg } +window_tuple = { + // empty tuple (0-tuple) + "(" ~ sp ~ ")" + // Single element (1-tuple) - uses trailing comma to distinguish from a bare value + | "(" ~ sp ~ window_value ~ sp ~ "," ~ sp ~ ")" + // Multi-element tuple. Trailing comma optional. + | "(" ~ sp ~ window_value ~ (sp ~ "," ~ sp ~ window_value)+ ~ (sp ~ ",")? ~ sp ~ ")" +} +window_value = { enum_value | untyped_literal | reference | expression | window_tuple | empty } +window_named_arg = { name ~ sp ~ "=" ~ sp ~ window_value } window_named_arg_list = { (window_named_arg ~ (sp ~ "," ~ sp ~ window_named_arg)*)? } window_function_call = { - function_call_prefix ~ sp ~ "over" ~ sp ~ "(" ~ sp ~ window_named_arg_list ~ sp ~ ")" ~ ":" ~ sp ~ type + function_reference ~ sp ~ argument_list ~ sp ~ "over" ~ sp ~ "(" ~ sp ~ window_named_arg_list ~ sp ~ ")" ~ ":" ~ sp ~ type } // Top-level Expression Rule diff --git a/src/parser/expressions.rs b/src/parser/expressions.rs index c1f10d01..58db130a 100644 --- a/src/parser/expressions.rs +++ b/src/parser/expressions.rs @@ -10,6 +10,7 @@ use substrait::proto::expression::{ WindowFunction, cast, reference_segment, }; use substrait::proto::function_argument::ArgType; +use substrait::proto::sort_field::{SortDirection, SortKind}; use substrait::proto::r#type::{Fp64, I64, Kind, Nullability}; use substrait::proto::{ AggregateFunction, AggregationPhase, Expression, FunctionArgument, SortField, Type, @@ -22,7 +23,7 @@ use super::{ }; use crate::extensions::SimpleExtensions; use crate::extensions::simple::{CompoundName, ExtensionKind}; -use crate::parser::relations::parse_expression_list; +use crate::parser::relations::ParsedNamedArgs; /// A field index (e.g., parsed from "$0" -> 0). #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -401,48 +402,92 @@ impl ScopedParsePair for Literal { } } -/// The shared prefix of function calls. -struct FunctionHead { +/// An unresolved reference to a function: its compound name plus an optional +/// explicit anchor, before it is looked up against the extension registry. +/// +/// This mirrors the `function_reference` grammar rule and is the shared +/// leading component of both scalar and window function calls. Resolution to a +/// concrete function anchor is a separate step ([`FunctionReference::resolve`]) +/// so that parsing stays free of extension lookups. +struct FunctionReference { name: CompoundName, anchor: Option, - arguments: Vec, } -/// Parse the `CompoundName ~ anchor? ~ urn_anchor? ~ argument_list` prefix of functions. -fn parse_function_head( - extensions: &SimpleExtensions, - iter: &mut RuleIter<'_>, -) -> Result { - // Parse compound function name (required) — e.g. "equal" or "equal:any_any" - let name = iter.parse_next::(); - - // Parse optional anchor (e.g., #1) - let anchor = iter - .try_pop(Rule::anchor) - .map(|n| unwrap_single_pair(n).as_str().parse::().unwrap()); - - // Parse optional URN anchor (e.g., @1) - let _urn_anchor = iter - .try_pop(Rule::urn_anchor) - .map(|n| unwrap_single_pair(n).as_str().parse::().unwrap()); - - // Parse argument list (required) - let argument_list = iter.pop(Rule::argument_list); - let mut arguments = Vec::new(); - for e in argument_list.into_inner() { - arguments.push(FunctionArgument { - arg_type: Some(ArgType::Value(Expression::parse_pair(extensions, e)?)), - }); +impl ParsePair for FunctionReference { + fn rule() -> Rule { + Rule::function_reference } - Ok(FunctionHead { - name, - anchor, - arguments, - }) + fn message() -> &'static str { + "FunctionReference" + } + + fn parse_pair(pair: pest::iterators::Pair) -> Self { + assert_eq!(pair.as_rule(), Self::rule()); + let mut iter = RuleIter::from(pair.into_inner()); + + // Compound function name (required) — e.g. "equal" or "equal:any_any" + let name = iter.parse_next::(); + + // Optional anchor (e.g., #1) + let anchor = iter + .try_pop(Rule::anchor) + .map(|n| unwrap_single_pair(n).as_str().parse::().unwrap()); + + // Optional URN anchor (e.g., @1); currently unused. + let _urn_anchor = iter + .try_pop(Rule::urn_anchor) + .map(|n| unwrap_single_pair(n).as_str().parse::().unwrap()); + + iter.done(); + FunctionReference { name, anchor } + } } -impl ScopedParsePair for ScalarFunction { +impl FunctionReference { + /// Resolve this reference to a concrete function anchor against the + /// extension registry. + fn resolve( + &self, + extensions: &SimpleExtensions, + span: pest::Span, + ) -> Result { + get_and_validate_anchor( + extensions, + ExtensionKind::Function, + self.anchor, + self.name.full(), + span, + ) + } +} + +/// Parse an `argument_list` rule (`(expr, expr, ...)`) into function arguments. +fn parse_argument_list( + extensions: &SimpleExtensions, + pair: pest::iterators::Pair, +) -> Result, MessageParseError> { + assert_eq!(pair.as_rule(), Rule::argument_list); + pair.into_inner() + .map(|e| { + Ok(FunctionArgument { + arg_type: Some(ArgType::Value(Expression::parse_pair(extensions, e)?)), + }) + }) + .collect() +} + +/// The parsed-but-unresolved form of a scalar function call: +/// `name#anchor(args):type`. The function reference is resolved against the +/// extension registry only in [`ScalarFunctionInvocation::resolve`]. +struct ScalarFunctionInvocation { + reference: FunctionReference, + arguments: Vec, + output_type: Type, +} + +impl ScopedParsePair for ScalarFunctionInvocation { fn rule() -> Rule { Rule::function_call } @@ -456,82 +501,324 @@ impl ScopedParsePair for ScalarFunction { pair: pest::iterators::Pair, ) -> Result { assert_eq!(pair.as_rule(), Self::rule()); - let span = pair.as_span(); let mut iter = RuleIter::from(pair.into_inner()); - let FunctionHead { - name, - anchor, - arguments, - } = parse_function_head(extensions, &mut iter)?; + // Drain the iterator into raw pairs before any fallible parsing, so an + // early return doesn't trip the RuleIter drop guard with pairs still + // pending. + let reference_pair = iter.pop(Rule::function_reference); + let args_pair = iter.pop(Rule::argument_list); + let type_pair = iter.pop(Rule::r#type); + iter.done(); - // Parse required output type (e.g., :i64). pop is safe here because - // the grammar guarantees the type token is always present. - let output_type = Some(Type::parse_pair(extensions, iter.pop(Rule::r#type))?); + let reference = FunctionReference::parse_pair(reference_pair); + let arguments = parse_argument_list(extensions, args_pair)?; + // Required output type (e.g., :i64); the grammar guarantees its presence. + let output_type = Type::parse_pair(extensions, type_pair)?; - iter.done(); - let anchor = get_and_validate_anchor( - extensions, - ExtensionKind::Function, - anchor, - name.full(), - span, - )?; - Ok(ScalarFunction { - function_reference: anchor, + Ok(ScalarFunctionInvocation { + reference, arguments, - options: vec![], // TODO: Function Options output_type, + }) + } +} + +impl ScalarFunctionInvocation { + /// Resolve the function reference and build the protobuf message. + fn resolve( + self, + extensions: &SimpleExtensions, + span: pest::Span, + ) -> Result { + let function_reference = self.reference.resolve(extensions, span)?; + Ok(ScalarFunction { + function_reference, + arguments: self.arguments, + options: vec![], // TODO: Function Options + output_type: Some(self.output_type), #[allow(deprecated)] args: vec![], }) } } -/// Parse a `window_bound` pair (`integer | empty`) into a `window_function::Bound`. -/// `_` (empty) means unbounded, represented explicitly as `Some(Bound { kind: -/// Some(Unbounded{}) })` rather than `None`. `0` maps to `CurrentRow`; -/// negative/positive integers map to `Preceding`/`Following`. -fn parse_window_bound( +impl ScopedParsePair for ScalarFunction { + fn rule() -> Rule { + Rule::function_call + } + + fn message() -> &'static str { + "ScalarFunction" + } + + fn parse_pair( + extensions: &SimpleExtensions, + pair: pest::iterators::Pair, + ) -> Result { + let span = pair.as_span(); + ScalarFunctionInvocation::parse_pair(extensions, pair)?.resolve(extensions, span) + } +} + +/// Resolve a `window_value` pair to an `Expression`, accepting either a bare +/// field reference (`$0`) or a general expression. +fn window_expression_from_value( + extensions: &SimpleExtensions, pair: pest::iterators::Pair, -) -> Result, MessageParseError> { - assert_eq!(pair.as_rule(), Rule::window_bound); +) -> Result { let inner = unwrap_single_pair(pair); match inner.as_rule() { - Rule::empty => Ok(Some(window_function::Bound { - kind: Some(bound::Kind::Unbounded(bound::Unbounded {})), - })), - Rule::integer => { - let offset: i64 = inner.as_str().parse().map_err(|_| { - MessageParseError::invalid( - "window_bound", - inner.as_span(), - format!("Window bound offset '{}' is out of range", inner.as_str()), - ) - })?; - let kind = match offset { - 0 => bound::Kind::CurrentRow(bound::CurrentRow {}), - n if n > 0 => bound::Kind::Following(bound::Following { offset: n }), - n => { - let offset = n.checked_neg().ok_or_else(|| { - MessageParseError::invalid( - "window_bound", - inner.as_span(), - format!("Window bound offset '{n}' is out of range"), - ) - })?; - bound::Kind::Preceding(bound::Preceding { offset }) + Rule::reference => Ok(Expression { + rex_type: Some(RexType::Selection(Box::new(FieldReference::parse_pair( + inner, + )))), + }), + Rule::expression => Expression::parse_pair(extensions, inner), + other => Err(MessageParseError::invalid( + "WindowFunction", + inner.as_span(), + format!("Expected an expression or field reference, got {other:?}"), + )), + } +} + +/// `partition=` accepts either one bare expression or a `window_tuple` of +/// expressions - a list-valued field written as a single item or a tuple. +fn parse_window_partition( + extensions: &SimpleExtensions, + pair: pest::iterators::Pair, +) -> Result, MessageParseError> { + let inner = unwrap_single_pair(pair.clone()); + match inner.as_rule() { + Rule::window_tuple => inner + .into_inner() + .map(|item| window_expression_from_value(extensions, item)) + .collect(), + _ => Ok(vec![window_expression_from_value(extensions, pair)?]), + } +} + +/// Build one `SortField` from a `window_tuple` with exactly two elements: a +/// field reference/expression and a sort-direction enum value. +fn sort_field_from_window_tuple( + extensions: &SimpleExtensions, + tuple: pest::iterators::Pair, +) -> Result { + assert_eq!(tuple.as_rule(), Rule::window_tuple); + let span = tuple.as_span(); + let items: Vec<_> = tuple.into_inner().collect(); + let [expr_pair, dir_pair] = <[_; 2]>::try_from(items).map_err(|items| { + MessageParseError::invalid( + "WindowFunction", + span, + format!( + "order= sort field must have exactly 2 elements (reference, direction), got {}", + items.len() + ), + ) + })?; + + let expr = window_expression_from_value(extensions, expr_pair)?; + + let dir_inner = unwrap_single_pair(dir_pair); + if dir_inner.as_rule() != Rule::enum_value { + return Err(MessageParseError::invalid( + "WindowFunction", + dir_inner.as_span(), + format!( + "order= sort field direction must be an enum value, got {:?}", + dir_inner.as_rule() + ), + )); + } + let direction = match dir_inner.as_str().trim_start_matches('&') { + "AscNullsFirst" => SortDirection::AscNullsFirst, + "AscNullsLast" => SortDirection::AscNullsLast, + "DescNullsFirst" => SortDirection::DescNullsFirst, + "DescNullsLast" => SortDirection::DescNullsLast, + other => { + return Err(MessageParseError::invalid( + "SortDirection", + dir_inner.as_span(), + format!("Unknown sort direction: {other}"), + )); + } + }; + Ok(SortField { + expr: Some(expr), + sort_kind: Some(SortKind::Direction(direction as i32)), + }) +} + +/// `order=` accepts one bare sort field (a `window_tuple` of `(reference, +/// direction)`) or a `window_tuple` of such sort fields, per the same +/// single-item-or-tuple convention as `partition=`. +fn parse_window_order( + extensions: &SimpleExtensions, + pair: pest::iterators::Pair, +) -> Result, MessageParseError> { + let inner = unwrap_single_pair(pair); + if inner.as_rule() != Rule::window_tuple { + return Err(MessageParseError::invalid( + "WindowFunction", + inner.as_span(), + format!( + "order= expects a sort field or a tuple of sort fields, got {:?}", + inner.as_rule() + ), + )); + } + + let items: Vec<_> = inner.clone().into_inner().collect(); + let is_list_of_fields = items + .first() + .map(|first| unwrap_single_pair(first.clone()).as_rule() == Rule::window_tuple) + .unwrap_or(false); + + if is_list_of_fields { + items + .into_iter() + .map(|item| sort_field_from_window_tuple(extensions, unwrap_single_pair(item))) + .collect() + } else { + Ok(vec![sort_field_from_window_tuple(extensions, inner)?]) + } +} + +impl ScopedParsePair for window_function::Bound { + fn rule() -> Rule { + Rule::window_value + } + + fn message() -> &'static str { + "WindowBound" + } + + /// A bound is always present once its enclosing `rows=`/`range=` frame is + /// present: `_` (empty) is an explicit `Unbounded`, not a missing value. + fn parse_pair( + _extensions: &SimpleExtensions, + pair: pest::iterators::Pair, + ) -> Result { + assert_eq!(pair.as_rule(), Self::rule()); + let inner = unwrap_single_pair(pair); + match inner.as_rule() { + Rule::empty => Ok(window_function::Bound { + kind: Some(bound::Kind::Unbounded(bound::Unbounded {})), + }), + Rule::untyped_literal => { + let lit = unwrap_single_pair(inner.clone()); + if lit.as_rule() != Rule::integer { + return Err(MessageParseError::invalid( + "WindowBound", + inner.as_span(), + format!( + "Window bound must be an integer or '_', got {:?}", + lit.as_rule() + ), + )); } - }; - Ok(Some(window_function::Bound { kind: Some(kind) })) + let offset: i64 = lit.as_str().parse().map_err(|_| { + MessageParseError::invalid( + "WindowBound", + lit.as_span(), + format!("Window bound offset '{}' is out of range", lit.as_str()), + ) + })?; + let kind = match offset { + 0 => bound::Kind::CurrentRow(bound::CurrentRow {}), + n if n > 0 => bound::Kind::Following(bound::Following { offset: n }), + n => { + let offset = n.checked_neg().ok_or_else(|| { + MessageParseError::invalid( + "WindowBound", + lit.as_span(), + format!("Window bound offset '{n}' is out of range"), + ) + })?; + bound::Kind::Preceding(bound::Preceding { offset }) + } + }; + Ok(window_function::Bound { kind: Some(kind) }) + } + other => Err(MessageParseError::invalid( + "WindowBound", + inner.as_span(), + format!("Window bound must be an integer or '_', got {other:?}"), + )), } - other => unreachable!("Grammar guarantees window_bound is integer or empty, got {other:?}"), } } -impl ScopedParsePair for WindowFunction { +/// `rows=`/`range=` accept a `window_tuple` of exactly two bounds (lower, upper). +fn parse_window_frame( + extensions: &SimpleExtensions, + pair: pest::iterators::Pair, +) -> Result<(window_function::Bound, window_function::Bound), MessageParseError> { + let inner = unwrap_single_pair(pair); + if inner.as_rule() != Rule::window_tuple { + return Err(MessageParseError::invalid( + "WindowFunction", + inner.as_span(), + format!( + "rows=/range= expects a tuple of two bounds, got {:?}", + inner.as_rule() + ), + )); + } + let span = inner.as_span(); + let items: Vec<_> = inner.into_inner().collect(); + let [lower_pair, upper_pair] = <[_; 2]>::try_from(items).map_err(|items| { + MessageParseError::invalid( + "WindowFunction", + span, + format!( + "rows=/range= expects exactly 2 bounds (lower, upper), got {}", + items.len() + ), + ) + })?; + let lower = window_function::Bound::parse_pair(extensions, lower_pair)?; + let upper = window_function::Bound::parse_pair(extensions, upper_pair)?; + Ok((lower, upper)) +} + +/// Resolve a `window_value` pair expected to hold an `enum_value` (e.g. +/// `&InitialToResult`), returning the identifier with its `&` prefix stripped. +fn window_enum_str( + pair: pest::iterators::Pair, + field: &'static str, +) -> Result { + let inner = unwrap_single_pair(pair); + if inner.as_rule() != Rule::enum_value { + return Err(MessageParseError::invalid( + "WindowFunction", + inner.as_span(), + format!("{field}= expects an enum value, got {:?}", inner.as_rule()), + )); + } + Ok(inner.as_str().trim_start_matches('&').to_string()) +} + +/// The parsed `over(...)` clause of a window function: partitioning, ordering, +/// aggregation phase/invocation, and an optional frame. All cross-field +/// validation (required `phase=`, mutually-exclusive `rows=`/`range=`, and the +/// single-`order=`-field requirement for `range=`) lives here, so the enclosing +/// [`WindowFunctionInvocation`] just composes components. +struct OverClause { + partitions: Vec, + sorts: Vec, + phase: i32, + invocation: i32, + bounds_type: i32, + lower_bound: Option, + upper_bound: Option, +} + +impl ScopedParsePair for OverClause { fn rule() -> Rule { - Rule::window_function_call + Rule::window_named_arg_list } fn message() -> &'static str { @@ -544,115 +831,92 @@ impl ScopedParsePair for WindowFunction { ) -> Result { assert_eq!(pair.as_rule(), Self::rule()); let span = pair.as_span(); - let mut iter = RuleIter::from(pair.into_inner()); - - let FunctionHead { - name, - anchor, - arguments, - } = parse_function_head(extensions, &mut iter)?; - // Parse the required `over(...)` named-argument list - let named_arg_list = iter.pop(Rule::window_named_arg_list); - - // Parse required output type (e.g., :i64) - let output_type = Some(Type::parse_pair(extensions, iter.pop(Rule::r#type))?); - iter.done(); - - let mut partitions = Vec::new(); - let mut sorts = Vec::new(); - let mut invocation = AggregationInvocation::Unspecified as i32; - let mut phase = None; - let mut bounds_type = window_function::BoundsType::Unspecified as i32; - let mut lower_bound = None; - let mut upper_bound = None; - - for arg in named_arg_list.into_inner() { - assert_eq!(arg.as_rule(), Rule::window_named_arg); - let inner = unwrap_single_pair(arg); - match inner.as_rule() { - Rule::window_partition_arg => { - let mut parts_iter = RuleIter::from(inner.into_inner()); - if let Some(expr_list) = parts_iter.try_pop(Rule::expression_list) { - partitions = parse_expression_list(extensions, expr_list)?; - } - parts_iter.done(); - } - Rule::window_order_arg => { - let mut order_iter = RuleIter::from(inner.into_inner()); - let sort_order = order_iter.pop(Rule::window_sort_order); - order_iter.done(); - for sf in sort_order.into_inner() { - sorts.push(SortField::parse_pair(extensions, sf)?); - } - } - Rule::window_frame_arg => { - let mut frame_iter = RuleIter::from(inner.into_inner()); - let kind = frame_iter.pop(Rule::window_frame_kind); - let lower = frame_iter.pop(Rule::window_bound); - let upper = frame_iter.pop(Rule::window_bound); - frame_iter.done(); - bounds_type = match kind.as_str() { - "rows" => window_function::BoundsType::Rows as i32, - "range" => window_function::BoundsType::Range as i32, - other => unreachable!( - "Grammar guarantees window_frame_kind is rows or range, got {other:?}" - ), - }; - lower_bound = parse_window_bound(lower)?; - upper_bound = parse_window_bound(upper)?; - } - Rule::window_phase_arg => { - let mut phase_iter = RuleIter::from(inner.into_inner()); - let enum_pair = phase_iter.pop(Rule::enum_value); - phase_iter.done(); - phase = Some(match enum_pair.as_str().trim_start_matches('&') { - "Unspecified" => AggregationPhase::Unspecified as i32, - "InitialToIntermediate" => AggregationPhase::InitialToIntermediate as i32, - "IntermediateToIntermediate" => { - AggregationPhase::IntermediateToIntermediate as i32 - } - "InitialToResult" => AggregationPhase::InitialToResult as i32, - "IntermediateToResult" => AggregationPhase::IntermediateToResult as i32, - other => { - return Err(MessageParseError::invalid( - "AggregationPhase", - enum_pair.as_span(), - format!("Unknown AggregationPhase: {other}"), - )); - } - }); - } - Rule::window_invocation_arg => { - let mut inv_iter = RuleIter::from(inner.into_inner()); - let enum_pair = inv_iter.pop(Rule::enum_value); - inv_iter.done(); - invocation = match enum_pair.as_str().trim_start_matches('&') { - "Unspecified" => AggregationInvocation::Unspecified as i32, - "All" => AggregationInvocation::All as i32, - "Distinct" => AggregationInvocation::Distinct as i32, - other => { - return Err(MessageParseError::invalid( - "AggregationInvocation", - enum_pair.as_span(), - format!("Unknown AggregationInvocation: {other}"), - )); - } - }; - } - other => { - unreachable!("Grammar guarantees window_named_arg alternatives, got {other:?}") - } - } + // `over(...)` reuses the same duplicate/unknown-argument-rejecting + // extractor used for Fetch's `limit=`/`offset=` (see + // `ParsedNamedArgs` in `src/parser/relations.rs`), instead of + // silently overwriting on duplicate names. + let extractor = ParsedNamedArgs::new(pair.into_inner(), Rule::window_named_arg)?; + let (extractor, partition_pair) = extractor.pop("partition", Rule::window_value); + let (extractor, order_pair) = extractor.pop("order", Rule::window_value); + let (extractor, phase_pair) = extractor.pop("phase", Rule::window_value); + let (extractor, invocation_pair) = extractor.pop("invocation", Rule::window_value); + let (extractor, rows_pair) = extractor.pop("rows", Rule::window_value); + let (extractor, range_pair) = extractor.pop("range", Rule::window_value); + extractor.done()?; + + if rows_pair.is_some() && range_pair.is_some() { + return Err(MessageParseError::invalid( + "WindowFunction", + span, + "rows= and range= are mutually exclusive", + )); } - let phase = phase.ok_or_else(|| { + let partitions = partition_pair + .map(|p| parse_window_partition(extensions, p)) + .transpose()? + .unwrap_or_default(); + let sorts = order_pair + .map(|p| parse_window_order(extensions, p)) + .transpose()? + .unwrap_or_default(); + + let phase_pair = phase_pair.ok_or_else(|| { MessageParseError::invalid( "WindowFunction", span, "Missing required phase= argument in over(...)", ) })?; + let phase = match window_enum_str(phase_pair, "phase")?.as_str() { + "Unspecified" => AggregationPhase::Unspecified as i32, + "InitialToIntermediate" => AggregationPhase::InitialToIntermediate as i32, + "IntermediateToIntermediate" => AggregationPhase::IntermediateToIntermediate as i32, + "InitialToResult" => AggregationPhase::InitialToResult as i32, + "IntermediateToResult" => AggregationPhase::IntermediateToResult as i32, + other => { + return Err(MessageParseError::invalid( + "AggregationPhase", + span, + format!("Unknown AggregationPhase: {other}"), + )); + } + }; + + let invocation = match invocation_pair { + None => AggregationInvocation::Unspecified as i32, + Some(p) => match window_enum_str(p, "invocation")?.as_str() { + "Unspecified" => AggregationInvocation::Unspecified as i32, + "All" => AggregationInvocation::All as i32, + "Distinct" => AggregationInvocation::Distinct as i32, + other => { + return Err(MessageParseError::invalid( + "AggregationInvocation", + span, + format!("Unknown AggregationInvocation: {other}"), + )); + } + }, + }; + + let (bounds_type, lower_bound, upper_bound) = if let Some(p) = rows_pair { + let (lower, upper) = parse_window_frame(extensions, p)?; + ( + window_function::BoundsType::Rows as i32, + Some(lower), + Some(upper), + ) + } else if let Some(p) = range_pair { + let (lower, upper) = parse_window_frame(extensions, p)?; + ( + window_function::BoundsType::Range as i32, + Some(lower), + Some(upper), + ) + } else { + (window_function::BoundsType::Unspecified as i32, None, None) + }; if bounds_type == window_function::BoundsType::Range as i32 && sorts.len() != 1 { return Err(MessageParseError::invalid( @@ -665,18 +929,90 @@ impl ScopedParsePair for WindowFunction { )); } - let anchor = get_and_validate_anchor( - extensions, - ExtensionKind::Function, - anchor, - name.full(), - span, - )?; - Ok(WindowFunction { - function_reference: anchor, + Ok(OverClause { + partitions, + sorts, + phase, + invocation, + bounds_type, + lower_bound, + upper_bound, + }) + } +} + +/// The parsed-but-unresolved form of a window function call: +/// `name#anchor(args) over(...):type`. The function reference is resolved +/// against the extension registry only in [`WindowFunctionInvocation::resolve`]. +struct WindowFunctionInvocation { + reference: FunctionReference, + arguments: Vec, + over: OverClause, + output_type: Type, +} + +impl ScopedParsePair for WindowFunctionInvocation { + fn rule() -> Rule { + Rule::window_function_call + } + + fn message() -> &'static str { + "WindowFunction" + } + + fn parse_pair( + extensions: &SimpleExtensions, + pair: pest::iterators::Pair, + ) -> Result { + assert_eq!(pair.as_rule(), Self::rule()); + let mut iter = RuleIter::from(pair.into_inner()); + + // Drain the iterator into raw pairs before any fallible parsing, so an + // early return doesn't trip the RuleIter drop guard with pairs still + // pending. + let reference_pair = iter.pop(Rule::function_reference); + let args_pair = iter.pop(Rule::argument_list); + let over_pair = iter.pop(Rule::window_named_arg_list); + let type_pair = iter.pop(Rule::r#type); + iter.done(); + + let reference = FunctionReference::parse_pair(reference_pair); + let arguments = parse_argument_list(extensions, args_pair)?; + let over = OverClause::parse_pair(extensions, over_pair)?; + // Required output type (e.g., :i64); the grammar guarantees its presence. + let output_type = Type::parse_pair(extensions, type_pair)?; + + Ok(WindowFunctionInvocation { + reference, arguments, - options: vec![], // TODO: Function Options + over, output_type, + }) + } +} + +impl WindowFunctionInvocation { + /// Resolve the function reference and build the protobuf message. + fn resolve( + self, + extensions: &SimpleExtensions, + span: pest::Span, + ) -> Result { + let function_reference = self.reference.resolve(extensions, span)?; + let OverClause { + partitions, + sorts, + phase, + invocation, + bounds_type, + lower_bound, + upper_bound, + } = self.over; + Ok(WindowFunction { + function_reference, + arguments: self.arguments, + options: vec![], // TODO: Function Options + output_type: Some(self.output_type), phase, sorts, invocation, @@ -690,6 +1026,24 @@ impl ScopedParsePair for WindowFunction { } } +impl ScopedParsePair for WindowFunction { + fn rule() -> Rule { + Rule::window_function_call + } + + fn message() -> &'static str { + "WindowFunction" + } + + fn parse_pair( + extensions: &SimpleExtensions, + pair: pest::iterators::Pair, + ) -> Result { + let span = pair.as_span(); + WindowFunctionInvocation::parse_pair(extensions, pair)?.resolve(extensions, span) + } +} + impl ScopedParsePair for Cast { fn rule() -> Rule { Rule::cast_expression @@ -1691,7 +2045,7 @@ mod tests { let exts = make_extensions_for_fn_tests(); let pair = parse_exact( Rule::window_function_call, - "add:i64_i64($0, $1) over(partition=($0), order=($1,&AscNullsLast), invocation=&Distinct, rows=(-3, 0), phase=&InitialToResult):i64", + "add:i64_i64($0, $1) over(partition=($0,), order=($1,&AscNullsLast), invocation=&Distinct, rows=(-3, 0), phase=&InitialToResult):i64", ); let f = WindowFunction::parse_pair(&exts, pair).unwrap(); assert_eq!(f.function_reference, 3); @@ -1779,7 +2133,7 @@ mod tests { let exts = make_extensions_for_fn_tests(); let pair = parse_exact( Rule::window_function_call, - "add:i64_i64($0, $1) over(partition=($0)):i64", + "add:i64_i64($0, $1) over(partition=($0,)):i64", ); let result = WindowFunction::parse_pair(&exts, pair); assert!(result.is_err(), "missing phase= must be rejected"); diff --git a/src/textify/expressions.rs b/src/textify/expressions.rs index b32549f7..6193aca9 100644 --- a/src/textify/expressions.rs +++ b/src/textify/expressions.rs @@ -19,8 +19,8 @@ use substrait::proto::{ use super::{PlanError, Scope, Textify, Visibility}; use crate::extensions::simple::ExtensionKind; -use crate::textify::rels::{Arguments, NamedArg, Value, ValueEnum, enum_str_value}; use crate::textify::types::{Name, NamedAnchor, OutputType, escaped}; +use crate::textify::values::{Arguments, NamedArg, Value, ValueEnum, enum_str_value}; // …(…) for function call // […] for variant @@ -541,20 +541,6 @@ impl Textify for IfThen { } } -fn window_enum_value<'a, T: TryFrom + ValueEnum>( - raw: i32, - field_name: &'static str, -) -> Value<'a> { - match T::try_from(raw) { - Ok(v) => enum_str_value(v.as_enum_str()), - Err(_) => Value::Missing(PlanError::invalid( - "WindowFunction", - Some(field_name), - format!("Unknown {field_name}: {raw}"), - )), - } -} - fn window_sort_order<'a>(sorts: &'a [SortField]) -> Value<'a> { match sorts { [single] => Value::from(single), @@ -562,24 +548,186 @@ fn window_sort_order<'a>(sorts: &'a [SortField]) -> Value<'a> { } } +/// A single partition expression is written bare (`partition=$0`) rather +/// than as a 1-tuple, since a bare parenthesized value (`($0)`) is not a +/// valid 1-tuple under the grammar - that form requires a trailing comma +/// (`($0,)`) to disambiguate from a parenthesized expression. +fn window_partition_value(partitions: &[Expression]) -> Value<'_> { + match partitions { + [single] => Value::Expression(single), + many => Value::Tuple(many.iter().map(Value::Expression).collect()), + } +} + +/// The frame keyword (`rows`/`range`) and bounds for a window function, if a +/// frame is present. `range=` additionally requires exactly one `order=` +/// field; a proto violating that invariant is still textified best-effort, +/// with a diagnostic surfaced via `ctx.push_error`. +fn window_frame_named_arg<'a, S: Scope>(f: &'a WindowFunction, ctx: &S) -> Option> { + let bounds_type = window_function::BoundsType::try_from(f.bounds_type); + let has_bounds = f.lower_bound.is_some() || f.upper_bound.is_some(); + if matches!(bounds_type, Ok(window_function::BoundsType::Unspecified)) && !has_bounds { + return None; + } + + let keyword = match bounds_type { + Ok(window_function::BoundsType::Rows) => "rows", + Ok(window_function::BoundsType::Range) => { + // The parser rejects range= frames unless there is exactly one + // order= field. A proto with a different count is still + // textified best-effort (the caller may have built it directly, + // bypassing the parser), but we surface a diagnostic since the + // output won't re-parse. + if f.sorts.len() != 1 { + ctx.push_error( + PlanError::invalid( + "WindowFunction", + Some("bounds_type"), + format!( + "range frame requires exactly one order= field, found {}", + f.sorts.len() + ), + ) + .into(), + ); + } + "range" + } + Ok(window_function::BoundsType::Unspecified) => { + // bounds_type is required whenever a frame is present. + ctx.push_error( + PlanError::invalid( + "WindowFunction", + Some("bounds_type"), + "bounds_type is Unspecified but lower_bound/upper_bound are set", + ) + .into(), + ); + "rows" + } + Err(_) => { + ctx.push_error( + PlanError::invalid( + "WindowFunction", + Some("bounds_type"), + format!("Unknown BoundsType: {}", f.bounds_type), + ) + .into(), + ); + "rows" + } + }; + let lower = window_bound_value(f.lower_bound.as_ref(), "lower_bound"); + let upper = window_bound_value(f.upper_bound.as_ref(), "upper_bound"); + Some(NamedArg { + name: Cow::Borrowed(keyword), + value: Value::Tuple(vec![lower, upper]), + }) +} + +/// Assembles the `over(...)` named arguments: `phase=`, `order=`, +/// `invocation=`, `partition=`, and `rows=`/`range=`. +fn window_over_named_args<'a, S: Scope>(f: &'a WindowFunction, ctx: &S) -> Vec> { + let mut named_args = Vec::new(); + + // phase= is always written, even when Unspecified: unlike invocation=, + // the parser requires a phase= argument to be present in over(...). + // + // The window textifier owns decoding of its raw enum fields: it decodes + // the i32, reports a field-specific error (tagged with the enum's own + // message type) on an unknown value, and otherwise hands the decoded + // value to the shared owned enum->Value conversion. + let phase_value = match AggregationPhase::try_from(f.phase) { + Ok(phase) => enum_str_value(phase.as_enum_str()), + Err(_) => Value::Missing(PlanError::invalid( + "AggregationPhase", + Some("phase"), + format!("Unknown AggregationPhase: {}", f.phase), + )), + }; + named_args.push(NamedArg { + name: Cow::Borrowed("phase"), + value: phase_value, + }); + + // order= is omitted when there are no sort fields. + if !f.sorts.is_empty() { + named_args.push(NamedArg { + name: Cow::Borrowed("order"), + value: window_sort_order(&f.sorts), + }); + } + + if f.invocation != AggregationInvocation::Unspecified as i32 { + let invocation_value = match AggregationInvocation::try_from(f.invocation) { + Ok(invocation) => enum_str_value(invocation.as_enum_str()), + Err(_) => Value::Missing(PlanError::invalid( + "AggregationInvocation", + Some("invocation"), + format!("Unknown AggregationInvocation: {}", f.invocation), + )), + }; + named_args.push(NamedArg { + name: Cow::Borrowed("invocation"), + value: invocation_value, + }); + } + + if !f.partitions.is_empty() { + named_args.push(NamedArg { + name: Cow::Borrowed("partition"), + value: window_partition_value(&f.partitions), + }); + } + + if let Some(frame) = window_frame_named_arg(f, ctx) { + named_args.push(frame); + } + + named_args +} + fn window_bound_value<'a>( bound: Option<&window_function::Bound>, field_name: &'static str, ) -> Value<'a> { - match bound.and_then(|b| b.kind.as_ref()) { - None | Some(bound::Kind::Unbounded(_)) => Value::EmptyGroup, + // `bound` being absent (no frame) and `bound.kind` being absent (a + // present-but-empty Bound, which the parser never produces) are distinct + // cases: the former is a legitimate unbounded frame, the latter is a + // malformed proto. Collapsing both into `EmptyGroup` would silently + // paper over the latter, so they're handled separately. + let Some(bound) = bound else { + return Value::EmptyGroup; + }; + match &bound.kind { + None => Value::Missing(PlanError::invalid( + "WindowFunction", + Some(field_name), + "Bound is present but has no kind set", + )), + Some(bound::Kind::Unbounded(_)) => Value::EmptyGroup, Some(bound::Kind::CurrentRow(_)) => Value::Integer(0), Some(bound::Kind::Preceding(p)) if p.offset < 0 => Value::Missing(PlanError::invalid( "WindowFunction", Some(field_name), format!("Preceding bound offset {} must not be negative", p.offset), )), + Some(bound::Kind::Preceding(p)) if p.offset == 0 => Value::Missing(PlanError::invalid( + "WindowFunction", + Some(field_name), + "Preceding bound offset must not be 0; use CurrentRow instead", + )), Some(bound::Kind::Preceding(p)) => Value::Integer(-p.offset), Some(bound::Kind::Following(f)) if f.offset < 0 => Value::Missing(PlanError::invalid( "WindowFunction", Some(field_name), format!("Following bound offset {} must not be negative", f.offset), )), + Some(bound::Kind::Following(f)) if f.offset == 0 => Value::Missing(PlanError::invalid( + "WindowFunction", + Some(field_name), + "Following bound offset must not be 0; use CurrentRow instead", + )), Some(bound::Kind::Following(f)) => Value::Integer(f.offset), } } @@ -590,6 +738,7 @@ impl Textify for WindowFunction { } fn textify(&self, ctx: &S, w: &mut W) -> fmt::Result { + // name/reference + arguments: `sum#10@1($0)` textify_function_call_prefix( self.function_reference, &self.arguments, @@ -597,99 +746,16 @@ impl Textify for WindowFunction { ctx, w, )?; - let output = OutputType(self.output_type.as_ref()); - let output_type = ctx.display(&output); - - let mut named_args: Vec = Vec::new(); - - // phase= is always written, even when Unspecified: unlike invocation=, - // the parser requires a phase= argument to be present in over(...). - named_args.push(NamedArg { - name: Cow::Borrowed("phase"), - value: window_enum_value::(self.phase, "phase"), - }); - - // order= is omitted when there are no sort fields. - if !self.sorts.is_empty() { - named_args.push(NamedArg { - name: Cow::Borrowed("order"), - value: window_sort_order(&self.sorts), - }); - } - - if self.invocation != AggregationInvocation::Unspecified as i32 { - named_args.push(NamedArg { - name: Cow::Borrowed("invocation"), - value: window_enum_value::(self.invocation, "invocation"), - }); - } - - if !self.partitions.is_empty() { - named_args.push(NamedArg { - name: Cow::Borrowed("partition"), - value: Value::Tuple(self.partitions.iter().map(Value::Expression).collect()), - }); - } - - let bounds_type = window_function::BoundsType::try_from(self.bounds_type); - let has_bounds = self.lower_bound.is_some() || self.upper_bound.is_some(); - if !matches!(bounds_type, Ok(window_function::BoundsType::Unspecified)) || has_bounds { - let keyword = match bounds_type { - Ok(window_function::BoundsType::Rows) => "rows", - Ok(window_function::BoundsType::Range) => { - // The parser rejects range= frames unless there is - // exactly one order= field; enforce the same invariant - // here so textified output always re-parses. - if self.sorts.len() != 1 { - ctx.push_error( - PlanError::invalid( - "WindowFunction", - Some("bounds_type"), - format!( - "range frame requires exactly one order= field, found {}", - self.sorts.len() - ), - ) - .into(), - ); - } - "range" - } - Ok(window_function::BoundsType::Unspecified) => { - // bounds_type is required whenever a frame is present. - ctx.push_error( - PlanError::invalid( - "WindowFunction", - Some("bounds_type"), - "bounds_type is Unspecified but lower_bound/upper_bound are set", - ) - .into(), - ); - "rows" - } - Err(_) => { - ctx.push_error( - PlanError::invalid( - "WindowFunction", - Some("bounds_type"), - format!("Unknown BoundsType: {}", self.bounds_type), - ) - .into(), - ); - "rows" - } - }; - let lower = window_bound_value(self.lower_bound.as_ref(), "lower_bound"); - let upper = window_bound_value(self.upper_bound.as_ref(), "upper_bound"); - named_args.push(NamedArg { - name: Cow::Borrowed(keyword), - value: Value::Tuple(vec![lower, upper]), - }); - } + // over-clause: ` over(phase=..., order=..., ...)` + let named_args = window_over_named_args(self, ctx); write!(w, " over(")?; Arguments::inline(vec![], named_args).textify(ctx, w)?; - write!(w, "){output_type}") + write!(w, ")")?; + + // output type: `:i64` + let output = OutputType(self.output_type.as_ref()); + write!(w, "{}", ctx.display(&output)) } } @@ -1351,7 +1417,7 @@ mod tests { }); assert_eq!( ctx.textify_no_errors(&f), - "sum() over(phase=&InitialToResult, order=($1, &AscNullsLast), invocation=&Distinct, partition=($0), rows=(-3, 0)):i16" + "sum() over(phase=&InitialToResult, order=($1, &AscNullsLast), invocation=&Distinct, partition=$0, rows=(-3, 0)):i16" ); } @@ -1457,6 +1523,54 @@ mod tests { assert!(!errs.is_empty(), "expected a diagnostic about the bound"); } + #[test] + fn test_window_function_bound_with_no_kind_surfaces_error() { + // A Bound message that is present but has no `kind` set is + // malformed - the parser never produces one. Collapsing it into the + // same `EmptyGroup` rendering as a genuinely absent bound would + // silently hide that the input was invalid. + let ctx = TestContext::new() + .with_urn(1, "urn:example") + .with_function(1, 10, "sum"); + let mut f = base_window_function(); + f.bounds_type = window_function::BoundsType::Rows as i32; + f.lower_bound = Some(window_function::Bound { kind: None }); + f.upper_bound = Some(window_function::Bound { + kind: Some(bound::Kind::CurrentRow(bound::CurrentRow {})), + }); + let (s, errs) = ctx.textify(&f); + assert_eq!( + s, + "sum() over(phase=&InitialToResult, rows=(!{WindowFunction}, 0)):i16" + ); + assert!(!errs.is_empty(), "expected a diagnostic about the bound"); + } + + #[test] + fn test_window_function_zero_offset_bound_surfaces_error() { + // A `Preceding`/`Following` offset of 0 is semantically equivalent + // to `CurrentRow`, but rendering it as a bare `0` would re-parse as + // `CurrentRow`, silently changing the bound's kind. Surface a + // diagnostic instead of rendering it. + let ctx = TestContext::new() + .with_urn(1, "urn:example") + .with_function(1, 10, "sum"); + let mut f = base_window_function(); + f.bounds_type = window_function::BoundsType::Rows as i32; + f.lower_bound = Some(window_function::Bound { + kind: Some(bound::Kind::Preceding(bound::Preceding { offset: 0 })), + }); + f.upper_bound = Some(window_function::Bound { + kind: Some(bound::Kind::Following(bound::Following { offset: 0 })), + }); + let (s, errs) = ctx.textify(&f); + assert_eq!( + s, + "sum() over(phase=&InitialToResult, rows=(!{WindowFunction}, !{WindowFunction})):i16" + ); + assert_eq!(errs.0.len(), 2, "expected a diagnostic for each bound"); + } + #[test] fn test_window_function_i64_max_preceding_bound_roundtrips() { // i64::MAX *does* have a negatable counterpart (-i64::MAX, which is @@ -1522,8 +1636,8 @@ mod tests { #[test] fn test_window_function_unknown_phase_and_invocation() { - // phase 99 and invocation 99 match no know n enum variant, so both - // render as the same generic `!{WindowFunction}` error token. + // phase 99 and invocation 99 match no known enum variant, so each + // renders as an error token tagged with that field's own enum type. let ctx = TestContext::new() .with_urn(1, "urn:example") .with_function(1, 10, "sum"); @@ -1533,7 +1647,7 @@ mod tests { let (s, errs) = ctx.textify(&f); assert_eq!( s, - "sum() over(phase=!{WindowFunction}, invocation=!{WindowFunction}):i16" + "sum() over(phase=!{AggregationPhase}, invocation=!{AggregationInvocation}):i16" ); assert!( !errs.is_empty(), diff --git a/src/textify/mod.rs b/src/textify/mod.rs index af541def..2bfde8ff 100644 --- a/src/textify/mod.rs +++ b/src/textify/mod.rs @@ -7,6 +7,7 @@ pub(crate) mod foundation; pub(crate) mod plan; pub(crate) mod rels; pub(crate) mod types; +pub(crate) mod values; #[cfg(test)] pub(crate) use foundation::ErrorQueue; diff --git a/src/textify/rels.rs b/src/textify/rels.rs index 07463575..d1278693 100644 --- a/src/textify/rels.rs +++ b/src/textify/rels.rs @@ -1,31 +1,26 @@ use std::borrow::Cow; use std::collections::HashMap; -use std::convert::TryFrom; use std::fmt; -use std::fmt::Debug; -use prost::{Message, UnknownEnumValue}; -use substrait::proto::aggregate_function::AggregationInvocation; +use prost::Message; use substrait::proto::fetch_rel::CountMode; use substrait::proto::plan_rel::RelType as PlanRelType; use substrait::proto::read_rel::ReadType; use substrait::proto::rel::RelType; use substrait::proto::rel_common::EmitKind; -use substrait::proto::sort_field::{SortDirection, SortKind}; use substrait::proto::{ - AggregateFunction, AggregateRel, AggregationPhase, CrossRel, Expression, ExtensionLeafRel, - ExtensionMultiRel, ExtensionSingleRel, FetchRel, FilterRel, JoinRel, NamedStruct, PlanRel, - ProjectRel, ReadRel, Rel, RelCommon, RelRoot, SetRel, SortField, SortRel, Type, join_rel, - set_rel, + AggregateRel, CrossRel, ExtensionLeafRel, ExtensionMultiRel, ExtensionSingleRel, FetchRel, + FilterRel, JoinRel, NamedStruct, PlanRel, ProjectRel, ReadRel, Rel, RelCommon, RelRoot, SetRel, + SortRel, join_rel, set_rel, }; use super::addenda::AddendumLines; -use super::expressions::Reference; use super::types::Name; +use super::values::{ArgsLayout, Arguments, NamedArg, Value, ValueEnum}; use super::{PlanError, Scope, Textify}; use crate::FormatError; use crate::extensions::any::AnyRef; -use crate::extensions::{ExtensionArgs, ExtensionColumn, ExtensionError, ExtensionValue}; +use crate::extensions::{ExtensionArgs, ExtensionError}; pub trait NamedRelation { fn name(&self) -> &'static str; @@ -75,82 +70,6 @@ impl Textify for Rel { /// Trait for enums that can be converted to a string representation for /// textification. -/// -/// Returns Ok(str) for valid enum values, or Err([PlanError]) for invalid or -/// unknown values. -pub trait ValueEnum { - fn as_enum_str(&self) -> Result, PlanError>; -} - -#[derive(Debug, Clone)] -pub struct NamedArg<'a> { - pub name: Cow<'a, str>, - pub value: Value<'a>, -} - -#[derive(Debug, Clone)] -pub enum Value<'a> { - TableName(Vec>), - Field(Option>, Option<&'a Type>), - Tuple(Vec>), - Reference(i32), - Expression(&'a Expression), - AggregateFunction(&'a AggregateFunction), - /// Represents a missing, invalid, or unspecified value. - Missing(PlanError), - /// Represents a valid enum value as a string for textification. - Enum(Cow<'a, str>), - EmptyGroup, - Integer(i64), - /// A decoded extension argument value. - ExtensionArgument(ExtensionValue), - /// A decoded extension output column. - ExtColumn(ExtensionColumn), -} - -impl<'a> Value<'a> { - pub fn expect(maybe_value: Option, f: impl FnOnce() -> PlanError) -> Self { - match maybe_value { - Some(s) => s, - None => Value::Missing(f()), - } - } -} - -impl<'a> From>, PlanError>> for Value<'a> { - fn from(token: Result>, PlanError>) -> Self { - match token { - Ok(value) => Value::TableName(value), - Err(err) => Value::Missing(err), - } - } -} - -impl<'a> Textify for Value<'a> { - fn name() -> &'static str { - "Value" - } - - fn textify(&self, ctx: &S, w: &mut W) -> fmt::Result { - match self { - Value::TableName(names) => write!(w, "{}", ctx.separated(names, ".")), - Value::Field(name, typ) => { - write!(w, "{}:{}", ctx.expect(name.as_ref()), ctx.expect(*typ)) - } - Value::Tuple(values) => write!(w, "({})", ctx.separated(values, ", ")), - Value::Reference(i) => write!(w, "{}", Reference(*i)), - Value::Expression(e) => write!(w, "{}", ctx.display(*e)), - Value::AggregateFunction(agg_fn) => agg_fn.textify(ctx, w), - Value::Missing(err) => write!(w, "{}", ctx.failure(err.clone())), - Value::Enum(res) => write!(w, "&{res}"), - Value::Integer(i) => write!(w, "{i}"), - Value::EmptyGroup => write!(w, "_"), - Value::ExtensionArgument(ev) => ev.textify(ctx, w), - Value::ExtColumn(ec) => ec.textify(ctx, w), - } - } -} - fn schema_to_values<'a>(schema: &'a NamedStruct) -> Vec> { let mut fields = schema .r#struct @@ -230,68 +149,6 @@ impl<'a> Textify for Emitted<'a> { } } -/// How an argument list renders inside a relation's `[...]`. -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -pub enum ArgsLayout { - /// `arg, arg, arg` on a single line. - #[default] - Inline, - /// One `- arg` per line, used for `Read:Virtual` rows. See - /// [`Relation::write_header`] for the exact layout. - Rows, -} - -#[derive(Debug, Clone)] -pub struct Arguments<'a> { - /// Positional arguments (e.g., a filter condition, group-bys, etc.) - pub positional: Vec>, - /// Named arguments (e.g., limit=10, offset=5) - pub named: Vec>, - /// How this argument list is laid out. Defaults to [`ArgsLayout::Inline`]; - /// only `Read:Virtual` opts into [`ArgsLayout::Rows`]. - layout: ArgsLayout, -} - -impl<'a> Arguments<'a> { - /// An inline argument list (`arg, arg, arg`), the default for every - /// relation. - pub fn inline(positional: Vec>, named: Vec>) -> Self { - Arguments { - positional, - named, - layout: ArgsLayout::Inline, - } - } - - /// A row-per-line argument list (`- arg` per line) used for `Read:Virtual` - /// with many rows. Currently not enabled for named arguments. - /// TODO: enable for named arguments as well. - pub fn rows(positional: Vec>) -> Self { - Arguments { - positional, - named: vec![], - layout: ArgsLayout::Rows, - } - } -} - -impl<'a> Textify for Arguments<'a> { - fn name() -> &'static str { - "Arguments" - } - fn textify(&self, ctx: &S, w: &mut W) -> fmt::Result { - if self.positional.is_empty() && self.named.is_empty() { - return write!(w, "_"); - } - - write!(w, "{}", ctx.separated(self.positional.iter(), ", "))?; - if !self.positional.is_empty() && !self.named.is_empty() { - write!(w, ", ")?; - } - write!(w, "{}", ctx.separated(self.named.iter(), ", ")) - } -} - pub struct Relation<'a> { pub name: Cow<'a, str>, /// Arguments to the relation, if any. @@ -355,7 +212,7 @@ impl Relation<'_> { None => { write!(w, "{indent}{name}[{cols}]") } - Some(args) if args.layout == ArgsLayout::Rows => { + Some(args) if args.layout() == ArgsLayout::Rows => { // One `- row` per line, one indent level deeper, with a trailing // comma on every row but the last, then `- => cols]`. let child = ctx.push_indent(); @@ -1100,178 +957,6 @@ impl<'a> Relation<'a> { } } -impl<'a> From<&'a SortField> for Value<'a> { - fn from(sf: &'a SortField) -> Self { - let field = match &sf.expr { - Some(expr) => match &expr.rex_type { - Some(substrait::proto::expression::RexType::Selection(fref)) => { - if let Some(substrait::proto::expression::field_reference::ReferenceType::DirectReference(seg)) = &fref.reference_type { - if let Some(substrait::proto::expression::reference_segment::ReferenceType::StructField(sf)) = &seg.reference_type { - Value::Reference(sf.field) - } else { Value::Missing(PlanError::unimplemented("SortField", Some("expr"), "Not a struct field")) } - } else { Value::Missing(PlanError::unimplemented("SortField", Some("expr"), "Not a direct reference")) } - } - _ => Value::Missing(PlanError::unimplemented( - "SortField", - Some("expr"), - "Not a selection", - )), - }, - None => Value::Missing(PlanError::unimplemented( - "SortField", - Some("expr"), - "Missing expr", - )), - }; - let direction = match &sf.sort_kind { - Some(kind) => Value::from(kind), - None => Value::Missing(PlanError::invalid( - "SortKind", - Some(Cow::Borrowed("sort_kind")), - "Missing sort_kind", - )), - }; - Value::Tuple(vec![field, direction]) - } -} - -/// Converts an [`ValueEnum::as_enum_str`] result into a [`Value`]. Shared by -/// the blanket `From<&T>` impl below and by callers that only have an owned -/// enum value (and so can't borrow it for the lifetime `Value<'a>` requires). -pub(crate) fn enum_str_value<'a>(result: Result, PlanError>) -> Value<'a> { - match result { - Ok(s) => Value::Enum(s), - Err(e) => Value::Missing(e), - } -} - -impl<'a, T: ValueEnum + ?Sized> From<&'a T> for Value<'a> { - fn from(enum_val: &'a T) -> Self { - enum_str_value(enum_val.as_enum_str()) - } -} - -impl ValueEnum for SortKind { - fn as_enum_str(&self) -> Result, PlanError> { - let d = match self { - &SortKind::Direction(d) => SortDirection::try_from(d), - SortKind::ComparisonFunctionReference(f) => { - return Err(PlanError::invalid( - "SortKind", - Some(Cow::Owned(format!("function reference{f}"))), - "SortKind::ComparisonFunctionReference unimplemented", - )); - } - }; - let s = match d { - Err(UnknownEnumValue(d)) => { - return Err(PlanError::invalid( - "SortKind", - Some(Cow::Owned(format!("unknown variant: {d:?}"))), - "Unknown SortDirection", - )); - } - Ok(SortDirection::AscNullsFirst) => "AscNullsFirst", - Ok(SortDirection::AscNullsLast) => "AscNullsLast", - Ok(SortDirection::DescNullsFirst) => "DescNullsFirst", - Ok(SortDirection::DescNullsLast) => "DescNullsLast", - Ok(SortDirection::Clustered) => "Clustered", - Ok(SortDirection::Unspecified) => { - return Err(PlanError::invalid( - "SortKind", - Option::>::None, - "Unspecified SortDirection", - )); - } - }; - Ok(Cow::Borrowed(s)) - } -} - -impl ValueEnum for join_rel::JoinType { - fn as_enum_str(&self) -> Result, PlanError> { - let s = match self { - join_rel::JoinType::Unspecified => { - return Err(PlanError::invalid( - "JoinType", - Option::>::None, - "Unspecified JoinType", - )); - } - join_rel::JoinType::Inner => "Inner", - join_rel::JoinType::Outer => "Outer", - join_rel::JoinType::Left => "Left", - join_rel::JoinType::Right => "Right", - join_rel::JoinType::LeftSemi => "LeftSemi", - join_rel::JoinType::RightSemi => "RightSemi", - join_rel::JoinType::LeftAnti => "LeftAnti", - join_rel::JoinType::RightAnti => "RightAnti", - join_rel::JoinType::LeftSingle => "LeftSingle", - join_rel::JoinType::RightSingle => "RightSingle", - join_rel::JoinType::LeftMark => "LeftMark", - join_rel::JoinType::RightMark => "RightMark", - }; - Ok(Cow::Borrowed(s)) - } -} - -impl ValueEnum for set_rel::SetOp { - fn as_enum_str(&self) -> Result, PlanError> { - let s = match self { - set_rel::SetOp::Unspecified => { - return Err(PlanError::invalid( - "SetOp", - Option::>::None, - "Unspecified SetOp", - )); - } - set_rel::SetOp::MinusPrimary => "MinusPrimary", - set_rel::SetOp::MinusPrimaryAll => "MinusPrimaryAll", - set_rel::SetOp::MinusMultiset => "MinusMultiset", - set_rel::SetOp::IntersectionPrimary => "IntersectionPrimary", - set_rel::SetOp::IntersectionMultiset => "IntersectionMultiset", - set_rel::SetOp::IntersectionMultisetAll => "IntersectionMultisetAll", - set_rel::SetOp::UnionDistinct => "UnionDistinct", - set_rel::SetOp::UnionAll => "UnionAll", - }; - Ok(Cow::Borrowed(s)) - } -} - -impl ValueEnum for AggregationPhase { - fn as_enum_str(&self) -> Result, PlanError> { - let s = match self { - AggregationPhase::Unspecified => "Unspecified", - AggregationPhase::InitialToIntermediate => "InitialToIntermediate", - AggregationPhase::IntermediateToIntermediate => "IntermediateToIntermediate", - AggregationPhase::InitialToResult => "InitialToResult", - AggregationPhase::IntermediateToResult => "IntermediateToResult", - }; - Ok(Cow::Borrowed(s)) - } -} - -impl ValueEnum for AggregationInvocation { - fn as_enum_str(&self) -> Result, PlanError> { - let s = match self { - AggregationInvocation::Unspecified => "Unspecified", - AggregationInvocation::All => "All", - AggregationInvocation::Distinct => "Distinct", - }; - Ok(Cow::Borrowed(s)) - } -} - -impl<'a> Textify for NamedArg<'a> { - fn name() -> &'static str { - "NamedArg" - } - fn textify(&self, ctx: &S, w: &mut W) -> fmt::Result { - write!(w, "{}=", self.name)?; - self.value.textify(ctx, w) - } -} - #[cfg(test)] mod tests { use substrait::proto::aggregate_rel::Grouping; @@ -1282,12 +967,13 @@ mod tests { use substrait::proto::rel_common::{Direct, Emit}; use substrait::proto::r#type::{self as ptype, Boolean, I64, Kind, Nullability, Struct}; use substrait::proto::{ - Expression, FunctionArgument, NamedStruct, ReadRel, Type, aggregate_rel, + AggregateFunction, Expression, FunctionArgument, NamedStruct, ReadRel, Type, aggregate_rel, }; use super::*; use crate::fixtures::TestContext; use crate::parser::expressions::FieldIndex; + use crate::textify::expressions::Reference; #[test] fn test_read_rel() { diff --git a/src/textify/values.rs b/src/textify/values.rs new file mode 100644 index 00000000..cd6f09e1 --- /dev/null +++ b/src/textify/values.rs @@ -0,0 +1,339 @@ +//! Shared value-rendering primitives ([`Value`], [`NamedArg`], [`Arguments`]) +//! used by both relation and expression textification. +//! +//! This module has no dependency on [`super::rels`] or [`super::expressions`] +//! beyond the tiny [`Reference`] display helper - relations and expressions +//! both build on top of it, rather than one depending on the other. + +use std::borrow::Cow; +use std::convert::TryFrom; +use std::fmt; + +use prost::UnknownEnumValue; +use substrait::proto::aggregate_function::AggregationInvocation; +use substrait::proto::sort_field::{SortDirection, SortKind}; +use substrait::proto::{ + AggregateFunction, AggregationPhase, Expression, SortField, join_rel, set_rel, +}; + +use super::expressions::Reference; +use super::types::Name; +use super::{PlanError, Scope, Textify}; +use crate::extensions::{ExtensionColumn, ExtensionValue}; + +/// A trait for enum types that can be rendered as `&VariantName` in the text +/// format. +/// +/// Returns Ok(str) for valid enum values, or Err([PlanError]) for invalid or +/// unknown values. +pub trait ValueEnum { + fn as_enum_str(&self) -> Result, PlanError>; +} + +#[derive(Debug, Clone)] +pub struct NamedArg<'a> { + pub name: Cow<'a, str>, + pub value: Value<'a>, +} + +#[derive(Debug, Clone)] +pub enum Value<'a> { + TableName(Vec>), + Field(Option>, Option<&'a substrait::proto::Type>), + Tuple(Vec>), + Reference(i32), + Expression(&'a Expression), + AggregateFunction(&'a AggregateFunction), + /// Represents a missing, invalid, or unspecified value. + Missing(PlanError), + /// Represents a valid enum value as a string for textification. + Enum(Cow<'a, str>), + EmptyGroup, + Integer(i64), + /// A decoded extension argument value. + ExtensionArgument(ExtensionValue), + /// A decoded extension output column. + ExtColumn(ExtensionColumn), +} + +impl<'a> Value<'a> { + pub fn expect(maybe_value: Option, f: impl FnOnce() -> PlanError) -> Self { + match maybe_value { + Some(s) => s, + None => Value::Missing(f()), + } + } +} + +impl<'a> From>, PlanError>> for Value<'a> { + fn from(token: Result>, PlanError>) -> Self { + match token { + Ok(value) => Value::TableName(value), + Err(err) => Value::Missing(err), + } + } +} + +impl<'a> Textify for Value<'a> { + fn name() -> &'static str { + "Value" + } + + fn textify(&self, ctx: &S, w: &mut W) -> fmt::Result { + match self { + Value::TableName(names) => write!(w, "{}", ctx.separated(names, ".")), + Value::Field(name, typ) => { + write!(w, "{}:{}", ctx.expect(name.as_ref()), ctx.expect(*typ)) + } + Value::Tuple(values) => write!(w, "({})", ctx.separated(values, ", ")), + Value::Reference(i) => write!(w, "{}", Reference(*i)), + Value::Expression(e) => write!(w, "{}", ctx.display(*e)), + Value::AggregateFunction(agg_fn) => agg_fn.textify(ctx, w), + Value::Missing(err) => write!(w, "{}", ctx.failure(err.clone())), + Value::Enum(res) => write!(w, "&{res}"), + Value::Integer(i) => write!(w, "{i}"), + Value::EmptyGroup => write!(w, "_"), + Value::ExtensionArgument(ev) => ev.textify(ctx, w), + Value::ExtColumn(ec) => ec.textify(ctx, w), + } + } +} + +/// How an argument list renders inside a relation's `[...]` or an +/// expression's `(...)`. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum ArgsLayout { + /// `arg, arg, arg` on a single line. + #[default] + Inline, + /// One `- arg` per line, used for `Read:Virtual` rows. See + /// [`super::rels::Relation::write_header`] for the exact layout. + Rows, +} + +#[derive(Debug, Clone)] +pub struct Arguments<'a> { + /// Positional arguments (e.g., a filter condition, group-bys, etc.) + pub positional: Vec>, + /// Named arguments (e.g., limit=10, offset=5) + pub named: Vec>, + /// How this argument list is laid out. Defaults to [`ArgsLayout::Inline`]; + /// only `Read:Virtual` opts into [`ArgsLayout::Rows`]. + layout: ArgsLayout, +} + +impl<'a> Arguments<'a> { + /// An inline argument list (`arg, arg, arg`), the default for every + /// relation. + pub fn inline(positional: Vec>, named: Vec>) -> Self { + Arguments { + positional, + named, + layout: ArgsLayout::Inline, + } + } + + /// A row-per-line argument list (`- arg` per line) used for `Read:Virtual` + /// with many rows. Currently not enabled for named arguments. + /// TODO: enable for named arguments as well. + pub fn rows(positional: Vec>) -> Self { + Arguments { + positional, + named: vec![], + layout: ArgsLayout::Rows, + } + } + + pub fn layout(&self) -> ArgsLayout { + self.layout + } +} + +impl<'a> Textify for Arguments<'a> { + fn name() -> &'static str { + "Arguments" + } + fn textify(&self, ctx: &S, w: &mut W) -> fmt::Result { + if self.positional.is_empty() && self.named.is_empty() { + return write!(w, "_"); + } + + write!(w, "{}", ctx.separated(self.positional.iter(), ", "))?; + if !self.positional.is_empty() && !self.named.is_empty() { + write!(w, ", ")?; + } + write!(w, "{}", ctx.separated(self.named.iter(), ", ")) + } +} + +impl<'a> From<&'a SortField> for Value<'a> { + fn from(sf: &'a SortField) -> Self { + let field = match &sf.expr { + Some(expr) => match &expr.rex_type { + Some(substrait::proto::expression::RexType::Selection(fref)) => { + if let Some(substrait::proto::expression::field_reference::ReferenceType::DirectReference(seg)) = &fref.reference_type { + if let Some(substrait::proto::expression::reference_segment::ReferenceType::StructField(sf)) = &seg.reference_type { + Value::Reference(sf.field) + } else { Value::Missing(PlanError::unimplemented("SortField", Some("expr"), "Not a struct field")) } + } else { Value::Missing(PlanError::unimplemented("SortField", Some("expr"), "Not a direct reference")) } + } + _ => Value::Missing(PlanError::unimplemented( + "SortField", + Some("expr"), + "Not a selection", + )), + }, + None => Value::Missing(PlanError::unimplemented( + "SortField", + Some("expr"), + "Missing expr", + )), + }; + let direction = match &sf.sort_kind { + Some(kind) => Value::from(kind), + None => Value::Missing(PlanError::invalid( + "SortKind", + Some(Cow::Borrowed("sort_kind")), + "Missing sort_kind", + )), + }; + Value::Tuple(vec![field, direction]) + } +} + +/// Converts an [`ValueEnum::as_enum_str`] result into a [`Value`]. Shared by +/// the blanket `From<&T>` impl below and by callers that only have an owned +/// enum value (and so can't borrow it for the lifetime `Value<'a>` requires). +pub(crate) fn enum_str_value<'a>(result: Result, PlanError>) -> Value<'a> { + match result { + Ok(s) => Value::Enum(s), + Err(e) => Value::Missing(e), + } +} + +impl<'a, T: ValueEnum + ?Sized> From<&'a T> for Value<'a> { + fn from(enum_val: &'a T) -> Self { + enum_str_value(enum_val.as_enum_str()) + } +} + +impl ValueEnum for SortKind { + fn as_enum_str(&self) -> Result, PlanError> { + let d = match self { + &SortKind::Direction(d) => SortDirection::try_from(d), + SortKind::ComparisonFunctionReference(f) => { + return Err(PlanError::invalid( + "SortKind", + Some(Cow::Owned(format!("function reference{f}"))), + "SortKind::ComparisonFunctionReference unimplemented", + )); + } + }; + let s = match d { + Err(UnknownEnumValue(d)) => { + return Err(PlanError::invalid( + "SortKind", + Some(Cow::Owned(format!("unknown variant: {d:?}"))), + "Unknown SortDirection", + )); + } + Ok(SortDirection::AscNullsFirst) => "AscNullsFirst", + Ok(SortDirection::AscNullsLast) => "AscNullsLast", + Ok(SortDirection::DescNullsFirst) => "DescNullsFirst", + Ok(SortDirection::DescNullsLast) => "DescNullsLast", + Ok(SortDirection::Clustered) => "Clustered", + Ok(SortDirection::Unspecified) => { + return Err(PlanError::invalid( + "SortKind", + Option::>::None, + "Unspecified SortDirection", + )); + } + }; + Ok(Cow::Borrowed(s)) + } +} + +impl ValueEnum for join_rel::JoinType { + fn as_enum_str(&self) -> Result, PlanError> { + let s = match self { + join_rel::JoinType::Unspecified => { + return Err(PlanError::invalid( + "JoinType", + Option::>::None, + "Unspecified JoinType", + )); + } + join_rel::JoinType::Inner => "Inner", + join_rel::JoinType::Outer => "Outer", + join_rel::JoinType::Left => "Left", + join_rel::JoinType::Right => "Right", + join_rel::JoinType::LeftSemi => "LeftSemi", + join_rel::JoinType::RightSemi => "RightSemi", + join_rel::JoinType::LeftAnti => "LeftAnti", + join_rel::JoinType::RightAnti => "RightAnti", + join_rel::JoinType::LeftSingle => "LeftSingle", + join_rel::JoinType::RightSingle => "RightSingle", + join_rel::JoinType::LeftMark => "LeftMark", + join_rel::JoinType::RightMark => "RightMark", + }; + Ok(Cow::Borrowed(s)) + } +} + +impl ValueEnum for set_rel::SetOp { + fn as_enum_str(&self) -> Result, PlanError> { + let s = match self { + set_rel::SetOp::Unspecified => { + return Err(PlanError::invalid( + "SetOp", + Option::>::None, + "Unspecified SetOp", + )); + } + set_rel::SetOp::MinusPrimary => "MinusPrimary", + set_rel::SetOp::MinusPrimaryAll => "MinusPrimaryAll", + set_rel::SetOp::MinusMultiset => "MinusMultiset", + set_rel::SetOp::IntersectionPrimary => "IntersectionPrimary", + set_rel::SetOp::IntersectionMultiset => "IntersectionMultiset", + set_rel::SetOp::IntersectionMultisetAll => "IntersectionMultisetAll", + set_rel::SetOp::UnionDistinct => "UnionDistinct", + set_rel::SetOp::UnionAll => "UnionAll", + }; + Ok(Cow::Borrowed(s)) + } +} + +impl ValueEnum for AggregationPhase { + fn as_enum_str(&self) -> Result, PlanError> { + let s = match self { + AggregationPhase::Unspecified => "Unspecified", + AggregationPhase::InitialToIntermediate => "InitialToIntermediate", + AggregationPhase::IntermediateToIntermediate => "IntermediateToIntermediate", + AggregationPhase::InitialToResult => "InitialToResult", + AggregationPhase::IntermediateToResult => "IntermediateToResult", + }; + Ok(Cow::Borrowed(s)) + } +} + +impl ValueEnum for AggregationInvocation { + fn as_enum_str(&self) -> Result, PlanError> { + let s = match self { + AggregationInvocation::Unspecified => "Unspecified", + AggregationInvocation::All => "All", + AggregationInvocation::Distinct => "Distinct", + }; + Ok(Cow::Borrowed(s)) + } +} + +impl<'a> Textify for NamedArg<'a> { + fn name() -> &'static str { + "NamedArg" + } + fn textify(&self, ctx: &S, w: &mut W) -> fmt::Result { + write!(w, "{}=", self.name)?; + self.value.textify(ctx, w) + } +} diff --git a/tests/plan_roundtrip.rs b/tests/plan_roundtrip.rs index 99ba367e..d3653eb5 100644 --- a/tests/plan_roundtrip.rs +++ b/tests/plan_roundtrip.rs @@ -1144,7 +1144,7 @@ Functions: === Plan Root[a, b, s] - Project[$0, $1, sum($1) over(phase=&InitialToResult, order=($1, &AscNullsLast), invocation=&Distinct, partition=($0), rows=(-3, 0)):fp64?] + Project[$0, $1, sum($1) over(phase=&InitialToResult, order=($1, &AscNullsLast), invocation=&Distinct, partition=$0, rows=(-3, 0)):fp64?] Read[t => a:i32, b:fp64]"#; roundtrip_plan(plan); @@ -1160,7 +1160,7 @@ Functions: === Plan Root[a, r] - Project[$0, row_number() over(phase=&InitialToResult, order=($0, &AscNullsLast), partition=($0)):i64] + Project[$0, row_number() over(phase=&InitialToResult, order=($0, &AscNullsLast), partition=$0):i64] Read[t => a:i32]"#; roundtrip_plan(plan); @@ -1250,7 +1250,7 @@ Functions: === Plan Root[r] - Project[row_number() over(partition=($0)):i64] + Project[row_number() over(partition=$0):i64] Read[t => a:i32]"#; assert!(Parser::parse(plan).is_err()); From 7efe5238b35578fc4cf869e993ae1b7c2e09f277 Mon Sep 17 00:00:00 2001 From: "gordon.hamilton" Date: Tue, 21 Jul 2026 16:18:36 -0400 Subject: [PATCH 6/7] feat: fixing code based on PR comments --- GRAMMAR.md | 10 +- src/parser/common.rs | 93 ++++++++++++ src/parser/expression_grammar.pest | 25 ++-- src/parser/expressions.rs | 226 +++++++++++------------------ src/parser/extensions.rs | 22 +++ src/parser/mod.rs | 4 +- src/parser/relations.rs | 85 ++--------- src/textify/expressions.rs | 33 ++--- src/textify/rels.rs | 16 +- src/textify/values.rs | 33 +++-- tests/plan_roundtrip.rs | 19 +++ 11 files changed, 283 insertions(+), 283 deletions(-) diff --git a/GRAMMAR.md b/GRAMMAR.md index 6b9226dd..6f953d47 100644 --- a/GRAMMAR.md +++ b/GRAMMAR.md @@ -507,7 +507,7 @@ A window function computes a value over a "window" of rows related to the curren #### Syntax -`window_function := function_signature anchor? urn_anchor? "(" (expression ("," expression)*)? ")" "over(" window_named_arguments? ")" ":" type` +`window_function := function_signature anchor? urn_anchor? "(" (expression ("," expression)*)? ")" "over(" named_arguments ")" ":" type` Rather than a bespoke grammar production per named argument, `over(...)` reuses the same generic [`named_arguments`](#arguments) production used elsewhere (`name "=" argument`, @@ -526,7 +526,7 @@ and what shape each value must have: each bound is an integer (negative = preceding, positive = following, `0` = current row) or `_` for unbounded -`partition=`, `order=`, `invocation=`, and `rows=`/`range=` are all optional and are omitted when empty; `phase=` is always required and always printed, though this (along with the mutual exclusion of `rows=`/`range=`, and the requirement that a `range=` frame have exactly one `order=` field) is enforced when parsing rather than by the grammar itself (the named-arg list as a whole is syntactically optional). +`partition=`, `order=`, `invocation=`, and `rows=`/`range=` are all optional and are omitted when empty; `phase=` is always required and always printed. The grammar itself only requires that at least one named argument be present; the specific rules — that `phase=` must appear, that `rows=`/`range=` are mutually exclusive, and that a `range=` frame have exactly one `order=` field — are all enforced by the parser rather than by the grammar. #### Examples @@ -616,7 +616,7 @@ Arguments in relations can be literals, expressions, enums, or tuples thereof. #### Syntax ```text -argument := enum / reference / literal / expression / tuple +argument := enum / reference / literal / expression / tuple / "_" tuple := "(" ")" // 0-tuple / "(" argument "," ")" // 1-tuple (trailing comma required) / "(" argument ("," argument)+ ","? ")" // 2+-tuple (trailing comma optional) @@ -624,6 +624,10 @@ arguments := argument ("," argument)* named_arguments := name "=" argument ("," name "=" argument)* ``` +`_` is accepted syntactically as an argument (it denotes an unbounded window frame +bound in `over(...)`), but is only meaningful there; consumers that don't support it — +extension-relation arguments, for example — reject it during parsing. + Tuples follow the Python/Rust trailing-comma convention to disambiguate from parenthesised expressions: `(x)` is a parenthesised expression, not a tuple. A trailing comma is required to form a 1-element tuple: `(x,)`. For 2+ elements the trailing comma is optional: `(x, y)` and diff --git a/src/parser/common.rs b/src/parser/common.rs index a6bdbda6..c5468217 100644 --- a/src/parser/common.rs +++ b/src/parser/common.rs @@ -1,6 +1,8 @@ +use std::collections::HashMap; use std::fmt; use pest_derive::Parser as PestDeriveParser; +use substrait::proto::sort_field::SortDirection; use thiserror::Error; use crate::extensions::SimpleExtensions; @@ -315,6 +317,97 @@ impl Drop for RuleIter<'_> { } } +/// A collection of named arguments (`name=value` pairs) extracted from a +/// named-argument-list rule, keyed by name with duplicate-name rejection. +/// +/// Shared by every consumer of the generic `name=value` grammar — the `Fetch` +/// relation's `limit=`/`offset=`, and window functions' `over(...)` arguments. +/// It lives here in `common` (rather than in `relations` or `expressions`) so +/// that neither of those modules has to depend on the other for it. +/// +/// The fluent API ensures all arguments are processed exactly once and none are +/// forgotten: [`pop`](Self::pop) consumes a known argument, and +/// [`done`](Self::done) errors on any that remain unconsumed. +pub(crate) struct ParsedNamedArgs<'a> { + map: HashMap<&'a str, pest::iterators::Pair<'a, Rule>>, +} + +impl<'a> ParsedNamedArgs<'a> { + pub(crate) fn new( + pairs: pest::iterators::Pairs<'a, Rule>, + rule: Rule, + ) -> Result { + let mut map = HashMap::new(); + for pair in pairs { + assert_eq!(pair.as_rule(), rule); + let mut inner = pair.clone().into_inner(); + let name_pair = inner.next().unwrap(); + let value_pair = inner.next().unwrap(); + assert_eq!(inner.next(), None); + let name = name_pair.as_str(); + if map.contains_key(name) { + return Err(MessageParseError::invalid( + "NamedArg", + name_pair.as_span(), + format!("Duplicate argument: {name}"), + )); + } + map.insert(name, value_pair); + } + Ok(Self { map }) + } + + // Returns the pair if it exists and matches the rule, otherwise None. + // Asserts that the rule must match the rule of the pair (and therefore + // panics in non-release-mode if not) + pub(crate) fn pop( + mut self, + name: &str, + rule: Rule, + ) -> (Self, Option>) { + let pair = self.map.remove(name).inspect(|pair| { + assert_eq!(pair.as_rule(), rule, "Rule mismatch for argument {name}"); + }); + (self, pair) + } + + // Returns an error if there are any unused arguments. + pub(crate) fn done(self) -> Result<(), MessageParseError> { + if let Some((name, pair)) = self.map.iter().next() { + return Err(MessageParseError::invalid( + "NamedArgExtractor", + // No span available for all unused args; use default. + pair.as_span(), + format!("Unknown argument: {name}"), + )); + } + Ok(()) + } +} + +/// Map a sort-direction enum identifier (without the leading `&`) to a +/// [`SortDirection`]. Shared by the `Sort` relation's `sort_field` parser and +/// the window function's `order=` parser, which reach it from different grammar +/// rules (`sort_direction` vs a generic `enum_value`) but accept the same set +/// of variant names. Lives in `common` so neither `relations` nor +/// `expressions` depends on the other for it. +pub(crate) fn sort_direction_from_str( + name: &str, + span: pest::Span, +) -> Result { + match name { + "AscNullsFirst" => Ok(SortDirection::AscNullsFirst), + "AscNullsLast" => Ok(SortDirection::AscNullsLast), + "DescNullsFirst" => Ok(SortDirection::DescNullsFirst), + "DescNullsLast" => Ok(SortDirection::DescNullsLast), + other => Err(MessageParseError::invalid( + "SortDirection", + span, + format!("Unknown sort direction: {other}"), + )), + } +} + #[cfg(test)] pub(crate) mod test_support { use pest::Parser as PestParser; diff --git a/src/parser/expression_grammar.pest b/src/parser/expression_grammar.pest index b074c698..d129d386 100644 --- a/src/parser/expression_grammar.pest +++ b/src/parser/expression_grammar.pest @@ -185,8 +185,8 @@ cast_expression = { "(" ~ sp ~ expression ~ sp ~ ")" ~ sp ~ "::" ~ sp ~ cast_fai // -- Window Function Calls -- // -// WindowFunction[Expression.WindowFunction]: a function_call followed -// by a "over(...)" clause carrying window-specific named arguments. +// WindowFunction[Expression.WindowFunction]: a function_reference and +// argument_list followed by an "over(...)" clause carrying named arguments. // Example: sum($0) over(partition=($1,$2), order=($3,&AscNullsLast), rows=(-3, 0), phase=&InitialToResult, invocation=&Distinct):fp64 // // - `partition=(...)` - optional list of partitioning expressions; a single @@ -201,20 +201,8 @@ cast_expression = { "(" ~ sp ~ expression ~ sp ~ ")" ~ sp ~ "::" ~ sp ~ cast_fai // for the current row) or `_` for unbounded/unspecified. // - `phase=&...` - required aggregation phase // - `invocation=&...` - optional invocation (e.g. `&All`, `&Distinct`) -window_tuple = { - // empty tuple (0-tuple) - "(" ~ sp ~ ")" - // Single element (1-tuple) - uses trailing comma to distinguish from a bare value - | "(" ~ sp ~ window_value ~ sp ~ "," ~ sp ~ ")" - // Multi-element tuple. Trailing comma optional. - | "(" ~ sp ~ window_value ~ (sp ~ "," ~ sp ~ window_value)+ ~ (sp ~ ",")? ~ sp ~ ")" -} -window_value = { enum_value | untyped_literal | reference | expression | window_tuple | empty } -window_named_arg = { name ~ sp ~ "=" ~ sp ~ window_value } -window_named_arg_list = { (window_named_arg ~ (sp ~ "," ~ sp ~ window_named_arg)*)? } - window_function_call = { - function_reference ~ sp ~ argument_list ~ sp ~ "over" ~ sp ~ "(" ~ sp ~ window_named_arg_list ~ sp ~ ")" ~ ":" ~ sp ~ type + function_reference ~ sp ~ argument_list ~ sp ~ "over" ~ sp ~ "(" ~ sp ~ extension_named_arguments ~ sp ~ ")" ~ ":" ~ sp ~ type } // Top-level Expression Rule @@ -421,7 +409,12 @@ tuple = { ~ (sp ~ "," ~ sp ~ extension_argument)+ ~ (sp ~ ",")? ~ sp ~ ")" } -extension_argument = { enum_value | untyped_literal | reference | expression | tuple } +// The generic argument-value grammar, shared by extension-relation arguments +// and window-function `over(...)` arguments. `empty` ("_") is syntactically +// valid everywhere; each consumer's Rust parser decides where it is +// semantically meaningful (window frame bounds) versus an error (everywhere +// else - see `ExtensionValue::parse_pair`). +extension_argument = { enum_value | untyped_literal | reference | expression | tuple | empty } // Named arguments (name=value pairs) extension_named_arguments = { extension_named_argument ~ (sp ~ "," ~ sp ~ extension_named_argument)* } diff --git a/src/parser/expressions.rs b/src/parser/expressions.rs index 58db130a..a5dfcfd8 100644 --- a/src/parser/expressions.rs +++ b/src/parser/expressions.rs @@ -10,7 +10,7 @@ use substrait::proto::expression::{ WindowFunction, cast, reference_segment, }; use substrait::proto::function_argument::ArgType; -use substrait::proto::sort_field::{SortDirection, SortKind}; +use substrait::proto::sort_field::SortKind; use substrait::proto::r#type::{Fp64, I64, Kind, Nullability}; use substrait::proto::{ AggregateFunction, AggregationPhase, Expression, FunctionArgument, SortField, Type, @@ -18,12 +18,11 @@ use substrait::proto::{ use super::types::get_and_validate_anchor; use super::{ - MessageParseError, ParsePair, Rule, RuleIter, ScopedParsePair, unescape_string, - unwrap_single_pair, + MessageParseError, ParsePair, ParsedNamedArgs, Rule, RuleIter, ScopedParsePair, + sort_direction_from_str, unescape_string, unwrap_single_pair, }; use crate::extensions::SimpleExtensions; use crate::extensions::simple::{CompoundName, ExtensionKind}; -use crate::parser::relations::ParsedNamedArgs; /// A field index (e.g., parsed from "$0" -> 0). #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -402,13 +401,7 @@ impl ScopedParsePair for Literal { } } -/// An unresolved reference to a function: its compound name plus an optional -/// explicit anchor, before it is looked up against the extension registry. -/// -/// This mirrors the `function_reference` grammar rule and is the shared -/// leading component of both scalar and window function calls. Resolution to a -/// concrete function anchor is a separate step ([`FunctionReference::resolve`]) -/// so that parsing stays free of extension lookups. +/// An unresolved reference to a function: its compound name plus an optional explicit anchor. struct FunctionReference { name: CompoundName, anchor: Option, @@ -478,16 +471,7 @@ fn parse_argument_list( .collect() } -/// The parsed-but-unresolved form of a scalar function call: -/// `name#anchor(args):type`. The function reference is resolved against the -/// extension registry only in [`ScalarFunctionInvocation::resolve`]. -struct ScalarFunctionInvocation { - reference: FunctionReference, - arguments: Vec, - output_type: Type, -} - -impl ScopedParsePair for ScalarFunctionInvocation { +impl ScopedParsePair for ScalarFunction { fn rule() -> Rule { Rule::function_call } @@ -501,6 +485,7 @@ impl ScopedParsePair for ScalarFunctionInvocation { pair: pest::iterators::Pair, ) -> Result { assert_eq!(pair.as_rule(), Self::rule()); + let span = pair.as_span(); let mut iter = RuleIter::from(pair.into_inner()); // Drain the iterator into raw pairs before any fallible parsing, so an @@ -516,52 +501,21 @@ impl ScopedParsePair for ScalarFunctionInvocation { // Required output type (e.g., :i64); the grammar guarantees its presence. let output_type = Type::parse_pair(extensions, type_pair)?; - Ok(ScalarFunctionInvocation { - reference, - arguments, - output_type, - }) - } -} - -impl ScalarFunctionInvocation { - /// Resolve the function reference and build the protobuf message. - fn resolve( - self, - extensions: &SimpleExtensions, - span: pest::Span, - ) -> Result { - let function_reference = self.reference.resolve(extensions, span)?; + // Resolve the function reference against the registry last, once the + // rest of the call has parsed cleanly. + let function_reference = reference.resolve(extensions, span)?; Ok(ScalarFunction { function_reference, - arguments: self.arguments, + arguments, options: vec![], // TODO: Function Options - output_type: Some(self.output_type), + output_type: Some(output_type), #[allow(deprecated)] args: vec![], }) } } -impl ScopedParsePair for ScalarFunction { - fn rule() -> Rule { - Rule::function_call - } - - fn message() -> &'static str { - "ScalarFunction" - } - - fn parse_pair( - extensions: &SimpleExtensions, - pair: pest::iterators::Pair, - ) -> Result { - let span = pair.as_span(); - ScalarFunctionInvocation::parse_pair(extensions, pair)?.resolve(extensions, span) - } -} - -/// Resolve a `window_value` pair to an `Expression`, accepting either a bare +/// Resolve a `argument` pair to an `Expression`, accepting either a bare /// field reference (`$0`) or a general expression. fn window_expression_from_value( extensions: &SimpleExtensions, @@ -583,15 +537,13 @@ fn window_expression_from_value( } } -/// `partition=` accepts either one bare expression or a `window_tuple` of -/// expressions - a list-valued field written as a single item or a tuple. fn parse_window_partition( extensions: &SimpleExtensions, pair: pest::iterators::Pair, ) -> Result, MessageParseError> { let inner = unwrap_single_pair(pair.clone()); match inner.as_rule() { - Rule::window_tuple => inner + Rule::tuple => inner .into_inner() .map(|item| window_expression_from_value(extensions, item)) .collect(), @@ -599,13 +551,11 @@ fn parse_window_partition( } } -/// Build one `SortField` from a `window_tuple` with exactly two elements: a -/// field reference/expression and a sort-direction enum value. -fn sort_field_from_window_tuple( +fn sort_field_from_tuple( extensions: &SimpleExtensions, tuple: pest::iterators::Pair, ) -> Result { - assert_eq!(tuple.as_rule(), Rule::window_tuple); + assert_eq!(tuple.as_rule(), Rule::tuple); let span = tuple.as_span(); let items: Vec<_> = tuple.into_inner().collect(); let [expr_pair, dir_pair] = <[_; 2]>::try_from(items).map_err(|items| { @@ -632,34 +582,20 @@ fn sort_field_from_window_tuple( ), )); } - let direction = match dir_inner.as_str().trim_start_matches('&') { - "AscNullsFirst" => SortDirection::AscNullsFirst, - "AscNullsLast" => SortDirection::AscNullsLast, - "DescNullsFirst" => SortDirection::DescNullsFirst, - "DescNullsLast" => SortDirection::DescNullsLast, - other => { - return Err(MessageParseError::invalid( - "SortDirection", - dir_inner.as_span(), - format!("Unknown sort direction: {other}"), - )); - } - }; + let direction = + sort_direction_from_str(dir_inner.as_str().trim_start_matches('&'), dir_inner.as_span())?; Ok(SortField { expr: Some(expr), sort_kind: Some(SortKind::Direction(direction as i32)), }) } -/// `order=` accepts one bare sort field (a `window_tuple` of `(reference, -/// direction)`) or a `window_tuple` of such sort fields, per the same -/// single-item-or-tuple convention as `partition=`. fn parse_window_order( extensions: &SimpleExtensions, pair: pest::iterators::Pair, ) -> Result, MessageParseError> { let inner = unwrap_single_pair(pair); - if inner.as_rule() != Rule::window_tuple { + if inner.as_rule() != Rule::tuple { return Err(MessageParseError::invalid( "WindowFunction", inner.as_span(), @@ -673,22 +609,22 @@ fn parse_window_order( let items: Vec<_> = inner.clone().into_inner().collect(); let is_list_of_fields = items .first() - .map(|first| unwrap_single_pair(first.clone()).as_rule() == Rule::window_tuple) + .map(|first| unwrap_single_pair(first.clone()).as_rule() == Rule::tuple) .unwrap_or(false); if is_list_of_fields { items .into_iter() - .map(|item| sort_field_from_window_tuple(extensions, unwrap_single_pair(item))) + .map(|item| sort_field_from_tuple(extensions, unwrap_single_pair(item))) .collect() } else { - Ok(vec![sort_field_from_window_tuple(extensions, inner)?]) + Ok(vec![sort_field_from_tuple(extensions, inner)?]) } } impl ScopedParsePair for window_function::Bound { fn rule() -> Rule { - Rule::window_value + Rule::extension_argument } fn message() -> &'static str { @@ -751,13 +687,12 @@ impl ScopedParsePair for window_function::Bound { } } -/// `rows=`/`range=` accept a `window_tuple` of exactly two bounds (lower, upper). fn parse_window_frame( extensions: &SimpleExtensions, pair: pest::iterators::Pair, ) -> Result<(window_function::Bound, window_function::Bound), MessageParseError> { let inner = unwrap_single_pair(pair); - if inner.as_rule() != Rule::window_tuple { + if inner.as_rule() != Rule::tuple { return Err(MessageParseError::invalid( "WindowFunction", inner.as_span(), @@ -784,8 +719,6 @@ fn parse_window_frame( Ok((lower, upper)) } -/// Resolve a `window_value` pair expected to hold an `enum_value` (e.g. -/// `&InitialToResult`), returning the identifier with its `&` prefix stripped. fn window_enum_str( pair: pest::iterators::Pair, field: &'static str, @@ -818,7 +751,7 @@ struct OverClause { impl ScopedParsePair for OverClause { fn rule() -> Rule { - Rule::window_named_arg_list + Rule::extension_named_arguments } fn message() -> &'static str { @@ -834,15 +767,15 @@ impl ScopedParsePair for OverClause { // `over(...)` reuses the same duplicate/unknown-argument-rejecting // extractor used for Fetch's `limit=`/`offset=` (see - // `ParsedNamedArgs` in `src/parser/relations.rs`), instead of + // `ParsedNamedArgs` in `src/parser/common.rs`), instead of // silently overwriting on duplicate names. - let extractor = ParsedNamedArgs::new(pair.into_inner(), Rule::window_named_arg)?; - let (extractor, partition_pair) = extractor.pop("partition", Rule::window_value); - let (extractor, order_pair) = extractor.pop("order", Rule::window_value); - let (extractor, phase_pair) = extractor.pop("phase", Rule::window_value); - let (extractor, invocation_pair) = extractor.pop("invocation", Rule::window_value); - let (extractor, rows_pair) = extractor.pop("rows", Rule::window_value); - let (extractor, range_pair) = extractor.pop("range", Rule::window_value); + let extractor = ParsedNamedArgs::new(pair.into_inner(), Rule::extension_named_argument)?; + let (extractor, partition_pair) = extractor.pop("partition", Rule::extension_argument); + let (extractor, order_pair) = extractor.pop("order", Rule::extension_argument); + let (extractor, phase_pair) = extractor.pop("phase", Rule::extension_argument); + let (extractor, invocation_pair) = extractor.pop("invocation", Rule::extension_argument); + let (extractor, rows_pair) = extractor.pop("rows", Rule::extension_argument); + let (extractor, range_pair) = extractor.pop("range", Rule::extension_argument); extractor.done()?; if rows_pair.is_some() && range_pair.is_some() { @@ -972,7 +905,7 @@ impl ScopedParsePair for WindowFunctionInvocation { // pending. let reference_pair = iter.pop(Rule::function_reference); let args_pair = iter.pop(Rule::argument_list); - let over_pair = iter.pop(Rule::window_named_arg_list); + let over_pair = iter.pop(Rule::extension_named_arguments); let type_pair = iter.pop(Rule::r#type); iter.done(); @@ -2111,46 +2044,61 @@ mod tests { } #[test] - fn test_window_function_bound_offset_too_many_digits_fails() { - // The grammar's `integer` rule permits arbitrarily many digits, so a - // literal wider than i64 range must be rejected by the `.parse()` - // call itself, not just by the separate i64::MIN-negation overflow - // path exercised elsewhere. + fn test_window_function_invalid_over_clauses_rejected() { + // Each over(...) input below is syntactically parseable but violates a + // semantic rule enforced in Rust, so must be rejected rather than + // silently accepted. Grouped as a table since every case is the same + // shape (parse -> expect error); the reason labels the failing case. let exts = make_extensions_for_fn_tests(); - let pair = parse_exact( - Rule::window_function_call, - "add:i64_i64($0, $1) over(rows=(-99999999999999999999, 0), phase=&InitialToResult):i64", - ); - let result = WindowFunction::parse_pair(&exts, pair); - assert!( - result.is_err(), - "a bound literal wider than i64 range must be rejected" - ); - } - - #[test] - fn test_window_function_missing_phase_fails() { - let exts = make_extensions_for_fn_tests(); - let pair = parse_exact( - Rule::window_function_call, - "add:i64_i64($0, $1) over(partition=($0,)):i64", - ); - let result = WindowFunction::parse_pair(&exts, pair); - assert!(result.is_err(), "missing phase= must be rejected"); - } - - #[test] - fn test_window_function_range_with_multiple_order_fields_fails() { - let exts = make_extensions_for_fn_tests(); - let pair = parse_exact( - Rule::window_function_call, - "add:i64_i64($0, $1) over(order=(($0,&AscNullsLast),($1,&AscNullsLast)), range=(_, 0), phase=&InitialToResult):i64", - ); - let result = WindowFunction::parse_pair(&exts, pair); - assert!( - result.is_err(), - "range= with more than one order= field must be rejected" - ); + let cases = [ + ( + "add:i64_i64($0, $1) over(partition=($0,)):i64", + "missing required phase= argument", + ), + ( + "add:i64_i64($0, $1) over(phase=&InitialToResult, phase=&InitialToResult):i64", + "duplicate named argument (ParsedNamedArgs must not silently overwrite)", + ), + ( + "add:i64_i64($0, $1) over(phase=&InitialToResult, bogus=$0):i64", + "unknown named argument (extractor done() completeness check)", + ), + ( + "add:i64_i64($0, $1) over(phase=&InitialToResult, rows=(0, 0), range=(0, 0)):i64", + "rows= and range= are mutually exclusive", + ), + ( + "add:i64_i64($0, $1) over(order=(($0,&AscNullsLast),($1,&AscNullsLast)), range=(_, 0), phase=&InitialToResult):i64", + "range= frame requires exactly one order= field", + ), + ( + "add:i64_i64($0, $1) over(phase=&Bogus):i64", + "unknown phase= enum value", + ), + ( + "add:i64_i64($0, $1) over(phase=&InitialToResult, invocation=&Bogus):i64", + "unknown invocation= enum value", + ), + ( + "add:i64_i64($0, $1) over(phase=&InitialToResult, rows=($0, 0)):i64", + "non-integer window bound (a reference is not a valid bound)", + ), + ( + "add:i64_i64($0, $1) over(phase=&InitialToResult, rows=(0,)):i64", + "window frame with != 2 bounds", + ), + // A literal wider than i64 range is rejected by the `.parse()` + // itself, distinct from the i64::MIN-negation overflow path. + ( + "add:i64_i64($0, $1) over(rows=(-99999999999999999999, 0), phase=&InitialToResult):i64", + "bound literal wider than i64 range", + ), + ]; + for (input, reason) in cases { + let pair = parse_exact(Rule::window_function_call, input); + let result = WindowFunction::parse_pair(&exts, pair); + assert!(result.is_err(), "should be rejected: {reason} -- input: {input}"); + } } #[test] diff --git a/src/parser/extensions.rs b/src/parser/extensions.rs index ebb621d8..3bc9f062 100644 --- a/src/parser/extensions.rs +++ b/src/parser/extensions.rs @@ -319,6 +319,16 @@ impl ScopedParsePair for ExtensionValue { let expr = Expression::parse_pair(extensions, inner)?; ExtensionValue::from(expr) } + // `_` (empty) is syntactically a valid `extension_argument` (it is meaningful + // as an unbounded window frame bound), but has no meaning as an + // extension-relation argument. Reject it explicitly. + Rule::empty => { + return Err(MessageParseError::invalid( + "ExtensionValue", + inner.as_span(), + "'_' is not a valid argument value here", + )); + } _ => panic!("Unexpected extension argument type: {:?}", inner.as_rule()), }) } @@ -624,6 +634,18 @@ mod tests { ExtensionValue::parse(&SimpleExtensions::default(), text).unwrap() } + #[test] + fn test_extension_value_empty_is_rejected() { + // `_` is a syntactically valid `extension_argument` (it is meaningful as an + // unbounded window frame bound), but has no meaning as an + // extension-relation argument. + let result = ExtensionValue::parse(&SimpleExtensions::default(), "_"); + assert!( + result.is_err(), + "expected `_` to be rejected as an extension value, got {result:?}" + ); + } + #[test] fn test_parse_urn_extension_declaration() { let line = "@1: /my/urn1"; diff --git a/src/parser/mod.rs b/src/parser/mod.rs index b5dbe9c8..72d62a6b 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -9,8 +9,8 @@ pub(crate) mod types; pub use common::MessageParseError; pub(crate) use common::{ - ErrorKind, ExpressionParser, ParsePair, Rule, RuleIter, ScopedParsePair, iter_pairs, - unescape_string, unwrap_single_pair, + ErrorKind, ExpressionParser, ParsePair, ParsedNamedArgs, Rule, RuleIter, ScopedParsePair, + iter_pairs, sort_direction_from_str, unescape_string, unwrap_single_pair, }; pub use errors::{ParseContext, ParseError, ParseResult}; pub use extensions::{ExpectedExtensionLine, ExtensionParseError}; diff --git a/src/parser/relations.rs b/src/parser/relations.rs index 5ceb8290..dd5f2bb3 100644 --- a/src/parser/relations.rs +++ b/src/parser/relations.rs @@ -9,14 +9,17 @@ 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::{SortDirection, SortKind}; +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, Rule, RuleIter, ScopedParsePair, unwrap_single_pair}; +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}; @@ -234,64 +237,6 @@ fn parse_emit(reference_list: Pair, direct_output_count: usize) -> (EmitKi (emit, output_count) } -/// Extracts named arguments from pest pairs with duplicate detection and completeness checking. -/// -/// Usage: `extractor.pop("limit", Rule::fetch_value).0.pop("offset", Rule::fetch_value).0.done()` -/// -/// The fluent API ensures all arguments are processed exactly once and none are forgotten. -pub struct ParsedNamedArgs<'a> { - map: HashMap<&'a str, Pair<'a, Rule>>, -} - -impl<'a> ParsedNamedArgs<'a> { - pub fn new( - pairs: pest::iterators::Pairs<'a, Rule>, - rule: Rule, - ) -> Result { - let mut map = HashMap::new(); - for pair in pairs { - assert_eq!(pair.as_rule(), rule); - let mut inner = pair.clone().into_inner(); - let name_pair = inner.next().unwrap(); - let value_pair = inner.next().unwrap(); - assert_eq!(inner.next(), None); - let name = name_pair.as_str(); - if map.contains_key(name) { - return Err(MessageParseError::invalid( - "NamedArg", - name_pair.as_span(), - format!("Duplicate argument: {name}"), - )); - } - map.insert(name, value_pair); - } - Ok(Self { map }) - } - - // Returns the pair if it exists and matches the rule, otherwise None. - // Asserts that the rule must match the rule of the pair (and therefore - // panics in non-release-mode if not) - pub fn pop(mut self, name: &str, rule: Rule) -> (Self, Option>) { - let pair = self.map.remove(name).inspect(|pair| { - assert_eq!(pair.as_rule(), rule, "Rule mismatch for argument {name}"); - }); - (self, pair) - } - - // Returns an error if there are any unused arguments. - pub fn done(self) -> Result<(), MessageParseError> { - if let Some((name, pair)) = self.map.iter().next() { - return Err(MessageParseError::invalid( - "NamedArgExtractor", - // No span available for all unused args; use default. - pair.as_span(), - format!("Unknown argument: {name}"), - )); - } - Ok(()) - } -} - impl RelationParsePair for ReadRel { fn rule() -> Rule { Rule::read_relation @@ -850,22 +795,10 @@ impl ScopedParsePair for SortField { let reference_pair = iter.pop(Rule::reference); let field_index = FieldIndex::parse_pair(reference_pair); let direction_pair = iter.pop(Rule::sort_direction); - // Strip the '&' prefix from enum syntax (e.g., "&AscNullsFirst" -> - // "AscNullsFirst") The grammar includes '&' to distinguish enums from - // identifiers, but the enum variant names don't include it - let direction = match direction_pair.as_str().trim_start_matches('&') { - "AscNullsFirst" => SortDirection::AscNullsFirst, - "AscNullsLast" => SortDirection::AscNullsLast, - "DescNullsFirst" => SortDirection::DescNullsFirst, - "DescNullsLast" => SortDirection::DescNullsLast, - other => { - return Err(MessageParseError::invalid( - "SortDirection", - direction_pair.as_span(), - format!("Unknown sort direction: {other}"), - )); - } - }; + let direction = sort_direction_from_str( + direction_pair.as_str().trim_start_matches('&'), + direction_pair.as_span(), + )?; iter.done(); Ok(SortField { expr: Some(Expression { diff --git a/src/textify/expressions.rs b/src/textify/expressions.rs index 6193aca9..b86e7bd5 100644 --- a/src/textify/expressions.rs +++ b/src/textify/expressions.rs @@ -20,7 +20,7 @@ use substrait::proto::{ use super::{PlanError, Scope, Textify, Visibility}; use crate::extensions::simple::ExtensionKind; use crate::textify::types::{Name, NamedAnchor, OutputType, escaped}; -use crate::textify::values::{Arguments, NamedArg, Value, ValueEnum, enum_str_value}; +use crate::textify::values::{Arguments, NamedArg, Value, decode_enum_field}; // …(…) for function call // […] for variant @@ -632,22 +632,11 @@ fn window_over_named_args<'a, S: Scope>(f: &'a WindowFunction, ctx: &S) -> VecValue conversion. - let phase_value = match AggregationPhase::try_from(f.phase) { - Ok(phase) => enum_str_value(phase.as_enum_str()), - Err(_) => Value::Missing(PlanError::invalid( - "AggregationPhase", - Some("phase"), - format!("Unknown AggregationPhase: {}", f.phase), - )), - }; + // `decode_enum_field` decodes the raw i32 and reports a field-specific + // error (tagged with the enum's own message type) on an unknown value. named_args.push(NamedArg { name: Cow::Borrowed("phase"), - value: phase_value, + value: decode_enum_field::(f.phase, "AggregationPhase", "phase"), }); // order= is omitted when there are no sort fields. @@ -659,17 +648,13 @@ fn window_over_named_args<'a, S: Scope>(f: &'a WindowFunction, ctx: &S) -> Vec enum_str_value(invocation.as_enum_str()), - Err(_) => Value::Missing(PlanError::invalid( - "AggregationInvocation", - Some("invocation"), - format!("Unknown AggregationInvocation: {}", f.invocation), - )), - }; named_args.push(NamedArg { name: Cow::Borrowed("invocation"), - value: invocation_value, + value: decode_enum_field::( + f.invocation, + "AggregationInvocation", + "invocation", + ), }); } diff --git a/src/textify/rels.rs b/src/textify/rels.rs index d1278693..86a5e733 100644 --- a/src/textify/rels.rs +++ b/src/textify/rels.rs @@ -16,7 +16,7 @@ use substrait::proto::{ use super::addenda::AddendumLines; use super::types::Name; -use super::values::{ArgsLayout, Arguments, NamedArg, Value, ValueEnum}; +use super::values::{ArgsLayout, Arguments, NamedArg, Value, ValueEnum, decode_enum_field}; use super::{PlanError, Scope, Textify}; use crate::FormatError; use crate::extensions::any::AnyRef; @@ -68,8 +68,6 @@ impl Textify for Rel { } } -/// Trait for enums that can be converted to a string representation for -/// textification. fn schema_to_values<'a>(schema: &'a NamedStruct) -> Vec> { let mut fields = schema .r#struct @@ -910,17 +908,7 @@ impl<'a> Relation<'a> { total_columns / children.len() }; - let op_value = match set_rel::SetOp::try_from(rel.op) { - Ok(op) => match op.as_enum_str() { - Ok(s) => Value::Enum(s), - Err(e) => Value::Missing(e), - }, - Err(_) => Value::Missing(PlanError::invalid( - "SetRel", - Some("op"), - format!("Unknown set op: {}", rel.op), - )), - }; + let op_value = decode_enum_field::(rel.op, "SetRel", "op"); let arguments = Some(Arguments::inline(vec![op_value], vec![])); let emit = get_emit(rel.common.as_ref()); diff --git a/src/textify/values.rs b/src/textify/values.rs index cd6f09e1..85833f37 100644 --- a/src/textify/values.rs +++ b/src/textify/values.rs @@ -1,9 +1,5 @@ //! Shared value-rendering primitives ([`Value`], [`NamedArg`], [`Arguments`]) //! used by both relation and expression textification. -//! -//! This module has no dependency on [`super::rels`] or [`super::expressions`] -//! beyond the tiny [`Reference`] display helper - relations and expressions -//! both build on top of it, rather than one depending on the other. use std::borrow::Cow; use std::convert::TryFrom; @@ -16,16 +12,12 @@ use substrait::proto::{ AggregateFunction, AggregationPhase, Expression, SortField, join_rel, set_rel, }; -use super::expressions::Reference; use super::types::Name; use super::{PlanError, Scope, Textify}; use crate::extensions::{ExtensionColumn, ExtensionValue}; /// A trait for enum types that can be rendered as `&VariantName` in the text /// format. -/// -/// Returns Ok(str) for valid enum values, or Err([PlanError]) for invalid or -/// unknown values. pub trait ValueEnum { fn as_enum_str(&self) -> Result, PlanError>; } @@ -86,7 +78,8 @@ impl<'a> Textify for Value<'a> { write!(w, "{}:{}", ctx.expect(name.as_ref()), ctx.expect(*typ)) } Value::Tuple(values) => write!(w, "({})", ctx.separated(values, ", ")), - Value::Reference(i) => write!(w, "{}", Reference(*i)), + // Field-reference syntax (`$N`); inlined rather than importing `expressions::Reference`. + Value::Reference(i) => write!(w, "${i}"), Value::Expression(e) => write!(w, "{}", ctx.display(*e)), Value::AggregateFunction(agg_fn) => agg_fn.textify(ctx, w), Value::Missing(err) => write!(w, "{}", ctx.failure(err.clone())), @@ -211,6 +204,28 @@ pub(crate) fn enum_str_value<'a>(result: Result, PlanError>) - } } +/// Decode a raw protobuf enum field (`i32`) into its shared [`Value`] +/// rendering: convert to the enum type and then to its `&Variant` string, or +/// produce a field-specific diagnostic when the raw value matches no variant. +/// +/// Shared by the callers that render an enum field straight from its `i32` +/// (window `phase=`/`invocation=`, `SetRel`'s `op`), so the +/// decode-or-diagnose shape lives in one place. `message` is the proto message +/// tag used for the failure token, `field` the offending field name. +pub(crate) fn decode_enum_field<'a, T>(raw: i32, message: &'static str, field: &'static str) -> Value<'a> +where + T: TryFrom + ValueEnum, +{ + match T::try_from(raw) { + Ok(v) => enum_str_value(v.as_enum_str()), + Err(_) => Value::Missing(PlanError::invalid( + message, + Some(field), + format!("Unknown {message}: {raw}"), + )), + } +} + impl<'a, T: ValueEnum + ?Sized> From<&'a T> for Value<'a> { fn from(enum_val: &'a T) -> Self { enum_str_value(enum_val.as_enum_str()) diff --git a/tests/plan_roundtrip.rs b/tests/plan_roundtrip.rs index d3653eb5..526ace8b 100644 --- a/tests/plan_roundtrip.rs +++ b/tests/plan_roundtrip.rs @@ -1294,6 +1294,25 @@ Root[r] roundtrip_plan(plan); } +/// A multi-field `order=` (a tuple of sort-field tuples) round-trips, +/// exercising the `is_list_of_fields` parse branch end-to-end rather than only +/// via rejection/textify-only tests. +#[test] +fn test_window_function_multi_order_field_roundtrip() { + let plan = r#"=== Extensions +URNs: + @ 1: https://github.com/substrait-io/substrait/blob/main/extensions/functions_arithmetic.yaml +Functions: + # 10 @ 1: row_number + +=== Plan +Root[a, b, r] + Project[$0, $1, row_number() over(phase=&InitialToResult, order=(($0, &AscNullsLast), ($1, &DescNullsFirst))):i64] + Read[t => a:i32, b:i32]"#; + + roundtrip_plan(plan); +} + /// `partition=(...)` with more than one expression round-trips, preserving /// both order and count. #[test] From 1bb4f88683f6bed5e92e1974a03fce5f9028de4a Mon Sep 17 00:00:00 2001 From: "gordon.hamilton" Date: Tue, 21 Jul 2026 16:23:59 -0400 Subject: [PATCH 7/7] feat: formatting --- src/parser/expressions.rs | 11 ++++++++--- src/parser/extensions.rs | 2 +- src/textify/values.rs | 6 +++++- 3 files changed, 14 insertions(+), 5 deletions(-) diff --git a/src/parser/expressions.rs b/src/parser/expressions.rs index a5dfcfd8..d1c77c25 100644 --- a/src/parser/expressions.rs +++ b/src/parser/expressions.rs @@ -582,8 +582,10 @@ fn sort_field_from_tuple( ), )); } - let direction = - sort_direction_from_str(dir_inner.as_str().trim_start_matches('&'), dir_inner.as_span())?; + let direction = sort_direction_from_str( + dir_inner.as_str().trim_start_matches('&'), + dir_inner.as_span(), + )?; Ok(SortField { expr: Some(expr), sort_kind: Some(SortKind::Direction(direction as i32)), @@ -2097,7 +2099,10 @@ mod tests { for (input, reason) in cases { let pair = parse_exact(Rule::window_function_call, input); let result = WindowFunction::parse_pair(&exts, pair); - assert!(result.is_err(), "should be rejected: {reason} -- input: {input}"); + assert!( + result.is_err(), + "should be rejected: {reason} -- input: {input}" + ); } } diff --git a/src/parser/extensions.rs b/src/parser/extensions.rs index 3bc9f062..f2727118 100644 --- a/src/parser/extensions.rs +++ b/src/parser/extensions.rs @@ -638,7 +638,7 @@ mod tests { fn test_extension_value_empty_is_rejected() { // `_` is a syntactically valid `extension_argument` (it is meaningful as an // unbounded window frame bound), but has no meaning as an - // extension-relation argument. + // extension-relation argument. let result = ExtensionValue::parse(&SimpleExtensions::default(), "_"); assert!( result.is_err(), diff --git a/src/textify/values.rs b/src/textify/values.rs index 85833f37..cccbe6bb 100644 --- a/src/textify/values.rs +++ b/src/textify/values.rs @@ -212,7 +212,11 @@ pub(crate) fn enum_str_value<'a>(result: Result, PlanError>) - /// (window `phase=`/`invocation=`, `SetRel`'s `op`), so the /// decode-or-diagnose shape lives in one place. `message` is the proto message /// tag used for the failure token, `field` the offending field name. -pub(crate) fn decode_enum_field<'a, T>(raw: i32, message: &'static str, field: &'static str) -> Value<'a> +pub(crate) fn decode_enum_field<'a, T>( + raw: i32, + message: &'static str, + field: &'static str, +) -> Value<'a> where T: TryFrom + ValueEnum, {