diff --git a/GRAMMAR.md b/GRAMMAR.md index 63635846..8e275566 100644 --- a/GRAMMAR.md +++ b/GRAMMAR.md @@ -1088,22 +1088,22 @@ Root[result] ### Sort Relation -The Sort relation specifies sort fields and directions for ordering the input: +The Sort relation specifies sort expressions and directions for ordering the input: -Sort[($0, &AscNullsFirst), ($1, &DescNullsLast) => $0, $1] +Sort[($0, &AscNullsFirst), (lower($1):string, &DescNullsLast) => $0, $1] #### Syntax ```text sort_relation := "Sort" "[" sort_fields "=>" reference_list "]" sort_fields := sort_field ("," sort_field)* -sort_field := "(" reference "," sort_direction ")" +sort_field := "(" expression "," sort_direction ")" sort_direction := "&AscNullsFirst" / "&AscNullsLast" / "&DescNullsFirst" / "&DescNullsLast" ``` #### Components -- Each sort field is a tuple: `(reference, sort_direction)` +- Each sort field is a tuple: `(expression, sort_direction)` - Sort directions follow the general `enum` syntax and specify null handling - `reference_list` - comma-separated list of field references to pass through diff --git a/src/parser/expression_grammar.pest b/src/parser/expression_grammar.pest index b6fd322d..3518570f 100644 --- a/src/parser/expression_grammar.pest +++ b/src/parser/expression_grammar.pest @@ -333,10 +333,10 @@ aggregate_output = { (expression ~ (sp ~ "," ~ sp ~ expression)*)? } // Empty grouping symbol empty = { "_" } -// SortRel: Sort[($0, &AscNullsFirst), ($2, &DescNullsLast) => ...] +// SortRel: Sort[($0, &AscNullsFirst), (lower($1):string, &DescNullsLast) => ...] sort_relation = { "Sort" ~ "[" ~ sort_field_list ~ sp ~ "=>" ~ sp ~ reference_list ~ "]" } sort_field_list = { (sort_field ~ (sp ~ "," ~ sp ~ sort_field)*)? } -sort_field = { "(" ~ sp ~ reference ~ sp ~ "," ~ sp ~ sort_direction ~ sp ~ ")" } +sort_field = { "(" ~ sp ~ expression ~ sp ~ "," ~ sp ~ sort_direction ~ sp ~ ")" } sort_direction = { "&AscNullsFirst" | "&AscNullsLast" | "&DescNullsFirst" | "&DescNullsLast" } // FetchRel: Fetch[limit=..., offset=... => ...] (named arguments only, any order, or _ for empty) @@ -443,4 +443,4 @@ duration_subseconds = { integer ~ subsecond_unit } // The sub-second units that correspond to a writable `interval_day` precision: // ms = 3, us = 6, ns = 9, ps = 12. -subsecond_unit = { "ms" | "us" | "ns" | "ps" } \ No newline at end of file +subsecond_unit = { "ms" | "us" | "ns" | "ps" } diff --git a/src/parser/relations.rs b/src/parser/relations.rs index ff9e1ed1..38d9c1c8 100644 --- a/src/parser/relations.rs +++ b/src/parser/relations.rs @@ -882,13 +882,13 @@ impl ScopedParsePair for SortField { } fn parse_pair( - _extensions: &SimpleExtensions, + extensions: &SimpleExtensions, pair: Pair, ) -> Result { assert_eq!(pair.as_rule(), Self::rule()); let mut iter = RuleIter::from(pair.into_inner()); - let reference_pair = iter.pop(Rule::reference); - let field_index = FieldIndex::parse_pair(reference_pair); + let expression_pair = iter.pop(Rule::expression); + let expression = Expression::parse_pair(extensions, expression_pair)?; let direction_pair = iter.pop(Rule::sort_direction); let direction = sort_direction_from_str( direction_pair.as_str().trim_start_matches('&'), @@ -896,11 +896,7 @@ impl ScopedParsePair for SortField { )?; iter.done(); Ok(SortField { - expr: Some(Expression { - rex_type: Some(RexType::Selection(Box::new( - field_index.to_field_reference(), - ))), - }), + expr: Some(expression), // TODO: Add support for SortKind::ComparisonFunctionReference sort_kind: Some(SortKind::Direction(direction as i32)), }) diff --git a/src/textify/values.rs b/src/textify/values.rs index 0e93906f..72bc2397 100644 --- a/src/textify/values.rs +++ b/src/textify/values.rs @@ -7,9 +7,6 @@ use std::fmt; use prost::UnknownEnumValue; use substrait::proto::aggregate_function::AggregationInvocation; -use substrait::proto::expression::RexType; -use substrait::proto::expression::field_reference::ReferenceType as FieldReferenceType; -use substrait::proto::expression::reference_segment::ReferenceType as SegmentReferenceType; use substrait::proto::sort_field::{SortDirection, SortKind}; use substrait::proto::{ AggregateFunction, AggregationPhase, Expression, SortField, Type, join_rel, set_rel, @@ -131,32 +128,7 @@ impl<'a> Textify for Arguments<'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(RexType::Selection(fref)) => { - if let Some(FieldReferenceType::DirectReference(seg)) = &fref.reference_type { - if let Some(SegmentReferenceType::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", - )), - }, + Some(expr) => Value::Expression(expr), None => Value::Missing(PlanError::unimplemented( "SortField", Some("expr"), diff --git a/tests/plan_roundtrip.rs b/tests/plan_roundtrip.rs index fe3a5b99..f038b8be 100644 --- a/tests/plan_roundtrip.rs +++ b/tests/plan_roundtrip.rs @@ -201,6 +201,22 @@ Root[a, b] roundtrip_plan(plan); } +#[test] +fn test_sort_relation_with_scalar_function_expressions_roundtrip() { + let plan = r#"=== Extensions +URNs: + @ 1: https://github.com/substrait-io/substrait/blob/main/extensions/functions_arithmetic.yaml +Functions: + # 10 @ 1: add + +=== Plan +Root[a, b] + Sort[(add($0, $1):i32, &AscNullsFirst), (add($1, 1:i32):i32, &DescNullsLast) => $0, $1] + Read[table => a:i32, b:i32]"#; + + roundtrip_plan(plan); +} + #[test] fn test_fetch_relation_roundtrip() { let plan_both = r#"=== Plan