feat: implementing window function expressions support - #185
Conversation
There was a problem hiding this comment.
Pull request overview
This PR adds end-to-end support for Substrait Expression.WindowFunction in substrait-explain, introducing a func(args) over(...) : type surface syntax and wiring it through the Pest grammar, Rust parser, and textifier so window functions can round-trip between text format and protobuf.
Changes:
- Extends the expression grammar to recognize
window_function_callwithover(...)named args (phase/partition/order/invocation/frame). - Implements parsing into
substrait::proto::expression::WindowFunctionand textification back to the new syntax. - Adds round-trip and unit tests plus user-facing documentation updates in
GRAMMAR.md.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/plan_roundtrip.rs | Adds integration round-trip coverage for window function plans (including unbounded bounds + invalid range/order combos). |
| src/textify/expressions.rs | Implements WindowFunction textification and hooks it into RexType::WindowFunction. |
| src/parser/expressions.rs | Implements WindowFunction parsing and allows window calls as general expressions. |
| src/parser/expression_grammar.pest | Adds Pest rules for window function calls and their over(...) argument forms. |
| GRAMMAR.md | Documents the new window function syntax and provides doctest examples. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| /// Write a single sort field as `($ref,&Direction)`, for use in `order=`. | ||
| fn textify_sort_field<S: Scope, W: fmt::Write>(sf: &SortField, ctx: &S, w: &mut W) -> fmt::Result { | ||
| let expr = ctx.expect(sf.expr.as_ref()); | ||
| write!(w, "({expr},")?; | ||
| match sort_direction_str(sf.sort_kind.as_ref()) { | ||
| Ok(s) => textify_enum(s, ctx, w)?, | ||
| Err(e) => write!(w, "{}", ctx.failure(e))?, | ||
| } | ||
| write!(w, ")") | ||
| } |
| let bounds_type = window_function::BoundsType::try_from(self.bounds_type); | ||
| let has_bounds = self.lower_bound.is_some() || self.upper_bound.is_some(); | ||
| if !matches!(bounds_type, Ok(window_function::BoundsType::Unspecified)) || has_bounds { | ||
| let keyword = match bounds_type { | ||
| Ok(window_function::BoundsType::Rows) => "rows", | ||
| Ok(window_function::BoundsType::Range) => "range", |
| let mut partitions = Vec::new(); | ||
| let mut sorts = Vec::new(); | ||
| let mut invocation = AggregationInvocation::Unspecified as i32; | ||
| let mut phase = None; | ||
| let mut bounds_type = window_function::BoundsType::Unspecified as i32; | ||
| let mut lower_bound = None; | ||
| let mut upper_bound = None; | ||
|
|
||
| for arg in named_arg_list.into_inner() { | ||
| assert_eq!(arg.as_rule(), Rule::window_named_arg); | ||
| let inner = unwrap_single_pair(arg); | ||
| match inner.as_rule() { |
There was a problem hiding this comment.
is this case worth addressing?
There was a problem hiding this comment.
Yes - we should probably either catch these and error, or at least a TODO comment, rather than silently ignore.
5a7c968 to
6326640
Compare
| let mut partitions = Vec::new(); | ||
| let mut sorts = Vec::new(); | ||
| let mut invocation = AggregationInvocation::Unspecified as i32; | ||
| let mut phase = None; | ||
| let mut bounds_type = window_function::BoundsType::Unspecified as i32; | ||
| let mut lower_bound = None; | ||
| let mut upper_bound = None; | ||
|
|
||
| for arg in named_arg_list.into_inner() { | ||
| assert_eq!(arg.as_rule(), Rule::window_named_arg); | ||
| let inner = unwrap_single_pair(arg); | ||
| match inner.as_rule() { |
| // The parser rejects range= frames unless there is | ||
| // exactly one order= field; enforce the same invariant | ||
| // here so textified output always re-parses. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 63266404a8
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| match bound.and_then(|b| b.kind.as_ref()) { | ||
| None | Some(bound::Kind::Unbounded(_)) => Value::EmptyGroup, | ||
| Some(bound::Kind::CurrentRow(_)) => Value::Integer(0), | ||
| Some(bound::Kind::Preceding(p)) if p.offset < 0 => Value::Missing(PlanError::invalid( |
There was a problem hiding this comment.
When a protobuf contains a malformed window bound such as Preceding { offset: 0 }, this falls through and emits 0, which reparses as CurrentRow and silently changes the bound kind. The Substrait proto comments require preceding/following offsets to be strictly positive and say to use CurrentRow for zero, so the textifier should surface zero the same way it reports negative offsets instead of producing lossy text.
Useful? React with 👍 / 👎.
wackywendell
left a comment
wackywendell
left a comment
There was a problem hiding this comment.
This is a good start for window-function support, especially the round-trip coverage. I think the parser will be easier to extend if window functions reuse the existing argument and tuple syntax, with Rust responsible for interpreting the named fields. On the formatting side, the expressions module should not depend on the rels module; it should only go the other way to make the separation of concerns clearer, and we should reorganize as needed to achieve that.
There are also a couple of conversion cases to tighten up: representable range frames should not be rejected during parsing, and malformed protobuf bounds should produce diagnostics rather than silently changing variants. I'd like another pass on those boundaries before approval, but the overall direction looks good.
| window_phase_arg = { "phase" ~ sp ~ "=" ~ sp ~ enum_value } | ||
| window_invocation_arg = { "invocation" ~ sp ~ "=" ~ sp ~ enum_value } | ||
|
|
||
| window_named_arg = { window_partition_arg | window_order_arg | window_frame_arg | window_phase_arg | window_invocation_arg } |
There was a problem hiding this comment.
Could we reuse the common argument value, tuple, and named-argument grammar here, rather than define a separate production for each window field? The current approach gives window functions a parallel argument syntax, even though Rust still needs to interpret the fields afterwards.
I would expect Pest to parse the same generic values we use elsewhere, with the window parser then deciding which names are allowed and what shape each value must have. That would also give us one place in Rust to reject duplicate fields, require phase, and make rows and range mutually exclusive.
Tuple syntax should work the same way here as it does elsewhere: (x,) is a one-element tuple and (x, y) is a two-element tuple. Where a window field represents a list, its conversion can accept either one element or a tuple. For example, order=($0, &AscNullsLast) is one sort field, while order=(($0, &AscNullsLast),) is an explicit one-element tuple of sort fields. Canonical output can continue to use the shorter form.
| } | ||
|
|
||
| /// The shared prefix of function calls. | ||
| struct FunctionHead { |
There was a problem hiding this comment.
I'm not sure FunctionHead is the right abstraction here. It includes the function arguments, but leaves the output type and window clause for the caller to parse. Its name describes where it appears in the input rather than a complete thing in the syntax, so callers still need to know which pieces remain in the iterator.
Could we instead parse an unresolved function reference, arguments, an over clause, and the output type as meaningful components? Those could compose into complete ScalarFunctionInvocation and WindowFunctionInvocation syntax types. Conversion from those types could then resolve the function reference and construct the protobuf message.
| } | ||
|
|
||
| /// Parse the `CompoundName ~ anchor? ~ urn_anchor? ~ argument_list` prefix of functions. | ||
| fn parse_function_head( |
There was a problem hiding this comment.
Could the syntax types above implement the existing pair-parsing traits? Right now parse_function_head is a standalone function, even though ParsePair and ScopedParsePair normally connect a grammar rule directly to the Rust type it produces.
Function arguments need extension context, so ScopedParsePair seems like the natural fit for the invocation types. That would keep the rule and its parsing behavior on the parsed type, instead of asking callers to coordinate an iterator with a separate helper.
|
|
||
| use super::{PlanError, Scope, Textify, Visibility}; | ||
| use crate::extensions::simple::ExtensionKind; | ||
| use crate::textify::rels::{Arguments, NamedArg, Value, ValueEnum, enum_str_value}; |
There was a problem hiding this comment.
Could we avoid making expression formatting depend on the relation formatter? Arguments, NamedArg, Value, and ValueEnum describe shared text-format concepts, but they currently live in textify::rels. This change also puts the window-function enum implementations there, which makes it less clear which module owns this behavior.
One option would be to move Value, NamedArg, and the related enum conversion into a neutral textify module; you could also (instead?) move window functions into their own module. I'm open to other arrangements, but let's try and make them fairly well-defined/separate so that relations can depend on expressions (and treat it as a 'black box'), and not vice versa.
| } | ||
| } | ||
|
|
||
| fn window_enum_value<'a, T: TryFrom<i32> + ValueEnum>( |
There was a problem hiding this comment.
This interface looks like a symptom of the module boundary above. window_enum_value is generic over the protobuf enum type, but hard-codes WindowFunction in its errors. It also combines protobuf decoding, diagnostics, and formatting, and its return lifetime is unrelated to any input because ValueEnum only provides a borrowed conversion.
Could we give decoded enum values a direct, owned conversion into the shared formatting value instead? The window textifier could then keep responsibility for decoding its raw fields and reporting field-specific errors. I would prefer that over renaming or further generalizing this helper.
| let mut partitions = Vec::new(); | ||
| let mut sorts = Vec::new(); | ||
| let mut invocation = AggregationInvocation::Unspecified as i32; | ||
| let mut phase = None; | ||
| let mut bounds_type = window_function::BoundsType::Unspecified as i32; | ||
| let mut lower_bound = None; | ||
| let mut upper_bound = None; | ||
|
|
||
| for arg in named_arg_list.into_inner() { | ||
| assert_eq!(arg.as_rule(), Rule::window_named_arg); | ||
| let inner = unwrap_single_pair(arg); | ||
| match inner.as_rule() { |
There was a problem hiding this comment.
Yes - we should probably either catch these and error, or at least a TODO comment, rather than silently ignore.
| /// Write the `name(args, options)` prefix shared by `ScalarFunction`, | ||
| /// `AggregateFunction`, and `WindowFunction` textification. | ||
| fn textify_function_call_prefix<S: Scope, W: fmt::Write>( |
There was a problem hiding this comment.
As with the parser - I think this makes more sense to break into components: function name/reference; arguments; over clause; output type. Implement Textify for each of these (perhaps by creating a newtype for each, as needed), then Textify for a FunctionCall and for a WindowFunctionCall can just make four calls to their elements.
|
|
||
| #### Syntax | ||
|
|
||
| `window_function := function_signature anchor? urn_anchor? "(" (expression ("," expression)*)? ")" "over(" (window_named_arg ("," window_named_arg)*)? ")" ":" type` |
There was a problem hiding this comment.
As in the Pest grammar, I would define the arguments in over(…) as named_arguments here, which are already defined.
Below, defining the specific arguments allowed in terms of a "signature" for each (e.g. partition is an expression, or tuple of expressions), using terms as defined above, would be good.
| } | ||
|
|
||
| #[test] | ||
| fn test_window_function_full() { |
There was a problem hiding this comment.
There are a lot of complicated tests here, and by their nature, they are a bit hard to follow; do we need all of them, or are they covered by the roundtrip tests?
|
A lot has changed in this last review so I tried to summarize the changes claude made: Grammar: reuse the generic argument grammar for
|
…ub.com:DataDog/substrait-explain into gordon.hamilton/window-function-implementation
Description
Adds parsing and textifying for Substrait's
Expression.WindowFunction, following the existing scalar/aggregate function syntax with anover(...)clause for window-specific args.Syntax:
func(args) over(phase=..., partition=(...), order=..., invocation=..., rows=(lo, hi) | range=(lo, hi)):typephase=(required) andinvocation=map toAggregationPhase/AggregationInvocation.partition=(...)— partitioning expressions.order=— sort field(s); bare for a single field, parenthesized list of tuples for multiple, matchingSort'ssort_field_listconvention.rows=(lo, hi)/range=(lo, hi)— frame bounds, where each side is an integer offset,0for current row, or_for unbounded (parsed to an explicitBound::Unbounded, not left asNone).range=is rejected when more than oneorder=field is present, matching the proto's semantics.Changes:
src/parser/expression_grammar.pest— new grammar rules forwindow_function_calland itsover(...)args.src/parser/expressions.rs—WindowFunction::parse_pair.src/textify/expressions.rs— corresponding textifier.GRAMMAR.md— documents the new syntax.tests/plan_roundtrip.rs— round-trip integration tests, alongside new unit tests in the parser/textify modules.Type of Change
Testing
Related Issues
Closes #158