feat: Support expressions inside Sort relations

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
 

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" }
+subsecond_unit = { "ms" | "us" | "ns" | "ps" }

src/parser/relations.rs

@@ -882,25 +882,21 @@ impl ScopedParsePair for SortField {
     }
 
     fn parse_pair(
-        _extensions: &SimpleExtensions,
+        extensions: &SimpleExtensions,
         pair: Pair<Rule>,
     ) -> Result<Self, MessageParseError> {
         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('&'),
             direction_pair.as_span(),
         )?;
         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)),
         })

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"),

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