diff --git a/GRAMMAR.md b/GRAMMAR.md index 82333a46..d593b114 100644 --- a/GRAMMAR.md +++ b/GRAMMAR.md @@ -491,6 +491,77 @@ 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(" 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. 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 + +```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. @@ -536,7 +607,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) @@ -544,6 +615,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 57f67690..4117f9bf 100644 --- a/src/parser/expression_grammar.pest +++ b/src/parser/expression_grammar.pest @@ -146,15 +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 = { - function_signature ~ sp ~ anchor? ~ sp ~ urn_anchor? ~ sp ~ argument_list ~ ":" ~ sp ~ type + function_reference ~ sp ~ argument_list ~ ":" ~ sp ~ type } if_clause = { @@ -179,10 +183,34 @@ 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_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 +// 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_function_call = { + function_reference ~ sp ~ argument_list ~ sp ~ "over" ~ sp ~ "(" ~ sp ~ extension_named_arguments ~ 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. @@ -388,7 +416,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 6a2963dd..5ee7171d 100644 --- a/src/parser/expressions.rs +++ b/src/parser/expressions.rs @@ -1,20 +1,25 @@ 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::sort_field::SortKind; 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::{ - 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}; @@ -396,68 +401,584 @@ impl ScopedParsePair for Literal { } } -impl ScopedParsePair for ScalarFunction { +/// An unresolved reference to a function: its compound name plus an optional explicit anchor. +struct FunctionReference { + name: CompoundName, + anchor: Option, +} + +impl ParsePair for FunctionReference { fn rule() -> Rule { - Rule::function_call + Rule::function_reference } fn message() -> &'static str { - "ScalarFunction" + "FunctionReference" } - fn parse_pair( - extensions: &SimpleExtensions, - pair: pest::iterators::Pair, - ) -> Result { + fn parse_pair(pair: pest::iterators::Pair) -> Self { 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. "equal" or "equal:any_any" + // Compound function name (required) — e.g. "equal" or "equal:any_any" let name = iter.parse_next::(); - // Parse optional anchor (e.g., #1) + // 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) + // 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()); - // 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 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))?); - iter.done(); - let anchor = get_and_validate_anchor( + FunctionReference { name, anchor } + } +} + +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, - anchor, - name.full(), + 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() +} + +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 { + 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 + // 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(); + + 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)?; + + // 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: anchor, + function_reference, arguments, options: vec![], // TODO: Function Options + output_type: Some(output_type), + #[allow(deprecated)] + args: vec![], + }) + } +} + +/// 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, + pair: pest::iterators::Pair, +) -> Result { + let inner = unwrap_single_pair(pair); + match inner.as_rule() { + 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:?}"), + )), + } +} + +fn parse_window_partition( + extensions: &SimpleExtensions, + pair: pest::iterators::Pair, +) -> Result, MessageParseError> { + let inner = unwrap_single_pair(pair.clone()); + match inner.as_rule() { + Rule::tuple => inner + .into_inner() + .map(|item| window_expression_from_value(extensions, item)) + .collect(), + _ => Ok(vec![window_expression_from_value(extensions, pair)?]), + } +} + +fn sort_field_from_tuple( + extensions: &SimpleExtensions, + tuple: pest::iterators::Pair, +) -> Result { + 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| { + 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 = 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)), + }) +} + +fn parse_window_order( + extensions: &SimpleExtensions, + pair: pest::iterators::Pair, +) -> Result, MessageParseError> { + let inner = unwrap_single_pair(pair); + if inner.as_rule() != Rule::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::tuple) + .unwrap_or(false); + + if is_list_of_fields { + items + .into_iter() + .map(|item| sort_field_from_tuple(extensions, unwrap_single_pair(item))) + .collect() + } else { + Ok(vec![sort_field_from_tuple(extensions, inner)?]) + } +} + +impl ScopedParsePair for window_function::Bound { + fn rule() -> Rule { + Rule::extension_argument + } + + 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() + ), + )); + } + 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:?}"), + )), + } + } +} + +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::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)) +} + +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::extension_named_arguments + } + + 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(); + + // `over(...)` reuses the same duplicate/unknown-argument-rejecting + // extractor used for Fetch's `limit=`/`offset=` (see + // `ParsedNamedArgs` in `src/parser/common.rs`), instead of + // silently overwriting on duplicate names. + 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() { + return Err(MessageParseError::invalid( + "WindowFunction", + span, + "rows= and range= are mutually exclusive", + )); + } + + 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( + "WindowFunction", + span, + format!( + "range= frame requires exactly one order= field, got {}", + sorts.len() + ), + )); + } + + 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::extension_named_arguments); + 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, + 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, + partitions, + bounds_type, + lower_bound, + upper_bound, #[allow(deprecated)] args: vec![], }) } } +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 @@ -526,6 +1047,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 +1068,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() ), } @@ -1152,8 +1678,6 @@ mod tests { assert_eq!(pairs.as_str(), "equal:any_any"); } - // ---- Tests for ScalarFunction parsing with compound names ---- - fn make_extensions_for_fn_tests() -> SimpleExtensions { let mut exts = SimpleExtensions::default(); exts.add_extension_urn("urn".to_string(), 1).unwrap(); @@ -1446,4 +1970,146 @@ 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_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 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] + 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/parser/extensions.rs b/src/parser/extensions.rs index ebb621d8..f2727118 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 8d476a4c..f80d0dd4 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}; @@ -276,64 +279,6 @@ fn parse_emit_suffix(suffix: Pair) -> Option<(EmitKind, usize)> { Some((EmitKind::Emit(Emit { output_mapping }), 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 @@ -927,22 +872,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 7d9aaa67..b86e7bd5 100644 --- a/src/textify/expressions.rs +++ b/src/textify/expressions.rs @@ -1,20 +1,26 @@ -use std::fmt::{self}; +use std::borrow::Cow; +use std::fmt; 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::{ - AggregateFunction, Expression, FunctionArgument, FunctionOption, expression as expr, + AggregateFunction, AggregationPhase, Expression, FunctionArgument, FunctionOption, SortField, + expression as expr, }; 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, decode_enum_field}; // …(…) for function call // […] for variant @@ -402,32 +408,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)) } } @@ -522,6 +541,209 @@ impl Textify for IfThen { } } +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()), + } +} + +/// 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(...). + // `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: decode_enum_field::(f.phase, "AggregationPhase", "phase"), + }); + + // 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 { + named_args.push(NamedArg { + name: Cow::Borrowed("invocation"), + value: decode_enum_field::( + f.invocation, + "AggregationInvocation", + "invocation", + ), + }); + } + + 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> { + // `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), + } +} + +impl Textify for WindowFunction { + fn name() -> &'static str { + "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, + &self.options, + ctx, + w, + )?; + + // 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: `:i64` + let output = OutputType(self.output_type.as_ref()); + write!(w, "{}", ctx.display(&output)) + } +} + impl Textify for RexType { fn name() -> &'static str { "RexType" @@ -532,15 +754,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, @@ -645,26 +859,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)) } } @@ -672,6 +876,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::*; @@ -1135,4 +1340,303 @@ 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"); + } + + #[test] + 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::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 {})), + }); + 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_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 + // 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 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"); + let mut f = base_window_function(); + f.phase = 99; + f.invocation = 99; + let (s, errs) = ctx.textify(&f); + assert_eq!( + s, + "sum() over(phase=!{AggregationPhase}, invocation=!{AggregationInvocation}):i16" + ); + assert!( + !errs.is_empty(), + "expected diagnostics about phase/invocation" + ); + } } 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 b6c0eb94..5adfc72e 100644 --- a/src/textify/rels.rs +++ b/src/textify/rels.rs @@ -2,28 +2,26 @@ use std::borrow::Cow; use std::collections::HashSet; use std::convert::TryFrom; use std::fmt; -use std::fmt::Debug; -use prost::{Message, UnknownEnumValue}; +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, 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, decode_enum_field}; 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; @@ -71,84 +69,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 @@ -295,68 +215,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. @@ -423,7 +281,7 @@ impl Relation<'_> { let cols = ctx.display(&cols); 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(); @@ -1155,17 +1013,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()); @@ -1204,147 +1052,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]) - } -} - -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), - } - } -} - -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<'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; @@ -1355,12 +1062,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..cccbe6bb --- /dev/null +++ b/src/textify/values.rs @@ -0,0 +1,358 @@ +//! Shared value-rendering primitives ([`Value`], [`NamedArg`], [`Arguments`]) +//! used by both relation and expression textification. + +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::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. +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, ", ")), + // 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())), + 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), + } +} + +/// 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()) + } +} + +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 9b39dd24..64066035 100644 --- a/tests/plan_roundtrip.rs +++ b/tests/plan_roundtrip.rs @@ -1151,3 +1151,235 @@ 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()); +} + +/// 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()); +} + +/// 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); +} + +/// 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] +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); +}