ESQL: allow sorting by expressions and not only regular fields (#107158)

* Support expressions in sort commands

docs/changelog/107158.yaml

@@ -0,0 +1,5 @@
+pr: 107158
+summary: "ESQL: allow sorting by expressions and not only regular fields"
+area: ES|QL
+type: feature
+issues: []

x-pack/plugin/esql/qa/testFixtures/src/main/resources/eval.csv-spec

@@ -424,3 +424,69 @@ emp_no:i
 -10002
 -10003
 ;
+
+sortExpression1#[skip:-8.13.99,reason:supported in 8.14]
+FROM employees
+| SORT emp_no + salary ASC
+| EVAL emp_no = -emp_no
+| LIMIT 10
+| EVAL sum = -emp_no + salary
+| KEEP emp_no, salary, sum
+;
+
+  emp_no:i     |   salary:i    |      sum:i      
+-10015         |25324          |35339          
+-10035         |25945          |35980          
+-10092         |25976          |36068          
+-10048         |26436          |36484          
+-10057         |27215          |37272          
+-10084         |28035          |38119          
+-10026         |28336          |38362          
+-10068         |28941          |39009          
+-10060         |29175          |39235          
+-10042         |30404          |40446
+;
+
+sortConcat1#[skip:-8.13.99,reason:supported in 8.14]
+from employees
+| sort concat(left(last_name, 1), left(first_name, 1)), salary desc
+| keep first_name, last_name, salary
+| eval ll = left(last_name, 1), lf = left(first_name, 1)
+| limit 10
+;
+
+ first_name:keyword | last_name:keyword | salary:integer|ll:keyword|lf:keyword     
+Mona                |Azuma              |46595          |A         |M             
+Satosi              |Awdeh              |50249          |A         |S             
+Brendon             |Bernini            |33370          |B         |B             
+Breannda            |Billingsley        |29175          |B         |B             
+Cristinel           |Bouloucos          |58715          |B         |C             
+Charlene            |Brattka            |28941          |B         |C             
+Margareta           |Bierman            |41933          |B         |M             
+Mokhtar             |Bernatsky          |38992          |B         |M             
+Parto               |Bamford            |61805          |B         |P             
+Premal              |Baek               |52833          |B         |P
+;
+
+sortConcat2#[skip:-8.13.99,reason:supported in 8.14]
+from employees
+| eval ln = last_name, fn = first_name, concat = concat(left(last_name, 1), left(first_name, 1))
+| sort concat(left(ln, 1), left(fn, 1)), salary desc
+| keep f*, l*, salary
+| eval c = concat(left(last_name, 1), left(first_name, 1))
+| drop *name, lan*
+| limit 10
+;
+
+     fn:keyword     |     ln:keyword    | salary:integer|       c:keyword     
+Mona                |Azuma              |46595          |AM             
+Satosi              |Awdeh              |50249          |AS             
+Brendon             |Bernini            |33370          |BB             
+Breannda            |Billingsley        |29175          |BB             
+Cristinel           |Bouloucos          |58715          |BC             
+Charlene            |Brattka            |28941          |BC             
+Margareta           |Bierman            |41933          |BM             
+Mokhtar             |Bernatsky          |38992          |BM             
+Parto               |Bamford            |61805          |BP             
+Premal              |Baek               |52833          |BP
+;

x-pack/plugin/esql/qa/testFixtures/src/main/resources/stats.csv-spec

@@ -1585,3 +1585,27 @@ c:l | k1:i | languages:i
 21  | 5    | 5
 10  | null | null
 ;
+
+minWithSortExpression1#[skip:-8.13.99,reason:supported in 8.14]
+FROM employees | STATS min = min(salary) by languages | SORT min + languages;
+
+     min:i     |   languages:i   
+25324          |5              
+25976          |1              
+26436          |3              
+27215          |4              
+29175          |2              
+28336          |null           
+;
+
+minWithSortExpression2#[skip:-8.13.99,reason:supported in 8.14]
+FROM employees | STATS min = min(salary) by languages | SORT min + CASE(languages == 5, 655, languages);
+
+     min:i     |   languages:i   
+25976          |1              
+25324          |5              
+26436          |3              
+27215          |4              
+29175          |2              
+28336          |null           
+;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizer.java

@@ -84,6 +84,7 @@
 import static java.util.Arrays.asList;
 import static java.util.Collections.singleton;
 import static org.elasticsearch.xpack.esql.expression.NamedExpressions.mergeOutputExpressions;
+import static org.elasticsearch.xpack.esql.optimizer.LogicalPlanOptimizer.SubstituteSurrogates.rawTemporaryName;
 import static org.elasticsearch.xpack.ql.expression.Expressions.asAttributes;
 import static org.elasticsearch.xpack.ql.optimizer.OptimizerRules.TransformDirection;
 import static org.elasticsearch.xpack.ql.optimizer.OptimizerRules.TransformDirection.DOWN;
@@ -125,7 +126,8 @@ protected static Batch<LogicalPlan> substitutions() {
             new ReplaceRegexMatch(),
             new ReplaceAliasingEvalWithProject(),
             new SkipQueryOnEmptyMappings(),
-            new SubstituteSpatialSurrogates()
+            new SubstituteSpatialSurrogates(),
+            new ReplaceOrderByExpressionWithEval()
             // new NormalizeAggregate(), - waits on https://github.com/elastic/elasticsearch/issues/100634
         );
     }
@@ -321,6 +323,35 @@ protected SpatialRelatesFunction rule(SpatialRelatesFunction function) {
         }
     }
 
+    static class ReplaceOrderByExpressionWithEval extends OptimizerRules.OptimizerRule<OrderBy> {
+        private static int counter = 0;
+
+        @Override
+        protected LogicalPlan rule(OrderBy orderBy) {
+            int size = orderBy.order().size();
+            List<Alias> evals = new ArrayList<>(size);
+            List<Order> newOrders = new ArrayList<>(size);
+
+            for (int i = 0; i < size; i++) {
+                var order = orderBy.order().get(i);
+                if (order.child() instanceof Attribute == false) {
+                    var name = rawTemporaryName("order_by", String.valueOf(i), String.valueOf(counter++));
+                    var eval = new Alias(order.child().source(), name, order.child());
+                    newOrders.add(order.replaceChildren(List.of(eval.toAttribute())));
+                    evals.add(eval);
+                } else {
+                    newOrders.add(order);
+                }
+            }
+            if (evals.isEmpty()) {
+                return orderBy;
+            } else {
+                var newOrderBy = new OrderBy(orderBy.source(), new Eval(orderBy.source(), orderBy.child(), evals), newOrders);
+                return new Project(orderBy.source(), newOrderBy, orderBy.output());
+            }
+        }
+    }
+
     static class ConvertStringToByteRef extends OptimizerRules.OptimizerExpressionRule<Literal> {
 
         ConvertStringToByteRef() {