Add optimisation rule to pull OrderBy above InlineJoin

Author: bpintea

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

@@ -1,7 +1,3 @@
-//
-// TODO: re-enable the commented tests once the Join functionality stabilizes
-//
-
 allFieldsReturned
 required_capability: inlinestats_v9
 
@@ -881,7 +877,6 @@ emp_no:integer | languages:integer | gender:keyword | max_lang:integer | y:keywo
          10014 |                 5 | null           | 5                | null
 ;
 
-// fails with AssertionError at org.elasticsearch.xpack.esql.plan.logical.Limit.writeTo(Limit.java:70)
 groupByMultipleRenamedColumns_AndOneExpression_Last
 required_capability: inlinestats_v9
 
@@ -905,7 +900,6 @@ emp_no:integer | languages:integer | gender:keyword|first_name:keyword|max_lang:
 10010          |4                  |null           |Duangkaew         |4               |null           |4              |D              
 ;
 
-// fails with AssertionError at org.elasticsearch.xpack.esql.plan.logical.Limit.writeTo(Limit.java:70)
 groupByMultipleRenamedColumns_AndTwoExpressions
 required_capability: inlinestats_v9
 
@@ -929,7 +923,6 @@ emp_no:integer | languages:integer | gender:keyword|first_name:keyword|max_lang:
 10010          |4                  |null           |Duangkaew         |4               |D              |null           |D              |4              
 ;
 
-// fails with AssertionError at org.elasticsearch.xpack.esql.plan.logical.Limit.writeTo(Limit.java:70)
 groupByMultipleRenamedColumns_AndMultipleRenames
 required_capability: inlinestats_v9
 
@@ -954,7 +947,6 @@ emp_no:integer | languages:integer | gender:keyword|    f:keyword     |max_lang:
 10010          |4                  |null           |Duangkaew         |4               |null           |4              |D              
 ;
 
-// fails with AssertionError at org.elasticsearch.xpack.esql.plan.logical.Limit.writeTo(Limit.java:70)
 groupByMultipleRenamedColumns_AndSameNameExpressionGroupingOverride
 required_capability: inlinestats_v9
 
@@ -1477,3 +1469,101 @@ ROW salary = 12300, emp_no = 5, gender = "F"
 emp_no:integer
 5
 ;
+
+sortBeforeInlinestats1
+required_capability: inlinestats_v9
+
+ROW salary = 12300, emp_no = 5, gender = "F"
+| EVAL salaryK = salary/1000
+| SORT salaryK DESC
+| INLINESTATS sum = SUM(salaryK) BY gender
+| KEEP emp_no
+;
+
+emp_no:integer
+5
+;
+
+sortBeforeInlinestats2
+required_capability: inlinestats_v9
+
+FROM employees
+| SORT emp_no
+| EVAL salaryK = salary/1000
+| INLINESTATS count = COUNT(*) BY salaryK
+| KEEP emp_no, still_hired, count
+| LIMIT 5
+;
+
+emp_no:integer |still_hired:boolean|count:long
+10001          |true               |1
+10002          |true               |3
+10003          |false              |2
+10004          |true               |2
+10005          |true               |1
+;
+
+// fails with: java.lang.AssertionError: expected no concrete indices without data node plan
+sortBeforeInlinestats3-Ignore
+required_capability: inlinestats_v9
+
+FROM employees
+| SORT languages DESC
+| EVAL salaryK = salary/1000
+| INLINESTATS count = COUNT(*) BY salaryK
+| SORT emp_no
+| INLINESTATS min = MIN(MV_COUNT(languages)) BY salaryK
+| KEEP emp_no, still_hired, count
+| LIMIT 5
+;
+
+emp_no:integer |still_hired:boolean|count:long
+10001          |true               |1
+10002          |true               |3
+10003          |false              |2
+10004          |true               |2
+10005          |true               |1
+;
+
+// same as `afterLookup`, swapped SORT position
+sortBeforeInlinestatsAndLookupJoin
+required_capability: inlinestats_v9
+required_capability: join_lookup_v12
+
+FROM airports
+| EVAL backup_scalerank = scalerank
+| RENAME scalerank AS language_code
+| SORT abbrev DESC
+| LOOKUP JOIN languages_lookup ON language_code
+| RENAME language_name as scalerank
+| DROP language_code
+| INLINESTATS count=COUNT(*) BY scalerank
+| KEEP abbrev, *scalerank
+| LIMIT 5
+;
+
+abbrev:keyword |backup_scalerank:integer| scalerank:keyword
+null           |8                       |null
+null           |8                       |null
+null           |8                       |null
+ZRH            |3                       |Spanish
+ZNZ            |4                       |German
+;
+
+// same as `shadowingAggregateByNextGrouping`, swapped SORT position
+sortBeforeInlinestats
+required_capability: inlinestats_v9
+
+FROM employees
+| KEEP gender, languages, emp_no, salary
+| SORT emp_no
+| INLINESTATS gender = count_distinct(gender) BY languages
+| INLINESTATS avg(salary) BY gender
+| LIMIT 3
+;
+
+emp_no:integer |salary:integer |languages:integer|avg(salary):double|gender:long
+10001          |57305          |2                |48248.55          |2
+10002          |56371          |5                |48248.55          |2
+10003          |61805          |4                |48248.55          |2
+;

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

@@ -34,6 +34,7 @@
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneRedundantOrderBy;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneRedundantSortClauses;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneUnusedIndexMode;
+import org.elasticsearch.xpack.esql.optimizer.rules.logical.PullUpOrderByBeforeInlineJoin;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownAndCombineFilters;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownAndCombineLimits;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownAndCombineOrderBy;
@@ -203,6 +204,7 @@ protected static Batch<LogicalPlan> operators(boolean local) {
             new PushDownAndCombineOrderBy(),
             new PruneRedundantOrderBy(),
             new PruneRedundantSortClauses(),
+            new PullUpOrderByBeforeInlineJoin(),
             new PruneLeftJoinOnNullMatchingField()
         );
     }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PullUpOrderByBeforeInlineJoin.java

@@ -0,0 +1,69 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+package org.elasticsearch.xpack.esql.optimizer.rules.logical;
+
+import org.elasticsearch.xpack.esql.plan.logical.Limit;
+import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
+import org.elasticsearch.xpack.esql.plan.logical.OrderBy;
+import org.elasticsearch.xpack.esql.plan.logical.SortAgnostic;
+import org.elasticsearch.xpack.esql.plan.logical.join.InlineJoin;
+
+/**
+ * Pulls "up" an {@link OrderBy} node that is not preceded by a {@link Limit}, but is preceded by an {@link InlineJoin}.
+ * The InlineJoin is {@link SortAgnostic}, so the OrderBy can be pulled up without affecting the semantics of the join.
+ * This is needed since otherwise the OrderBy would remain to be executed unbounded, which isn't supported.
+ * If it's preceded by a {@link Limit}, it will be merged into a {@link org.elasticsearch.xpack.esql.plan.logical.TopN} later in the
+ * "cleanup" optimization stage.
+ */
+public final class PullUpOrderByBeforeInlineJoin extends OptimizerRules.OptimizerRule<LogicalPlan> {
+
+    @Override
+    protected LogicalPlan rule(LogicalPlan plan) {
+        return plan.transformUp(LogicalPlan.class, PullUpOrderByBeforeInlineJoin::pullUpOrderByBeforeInlineJoin);
+    }
+
+    private static LogicalPlan pullUpOrderByBeforeInlineJoin(LogicalPlan plan) {
+        if (plan instanceof InlineJoin inlineJoin) {
+            OrderBy orderBy = findOrderByNotPrecededByLimit(inlineJoin);
+            if (orderBy != null) {
+                LogicalPlan newInlineJoin = removeOrderBy(inlineJoin, orderBy);
+                return new OrderBy(orderBy.source(), newInlineJoin, orderBy.order());
+            }
+        }
+        return plan;
+    }
+
+    // Finds an OrderBy node in the subtree of the provided plan that is not preceded by a Limit
+    private static OrderBy findOrderByNotPrecededByLimit(LogicalPlan plan) {
+        if (plan instanceof Limit) {
+            return null;
+        }
+        if (plan instanceof OrderBy orderBy) {
+            return orderBy;
+        }
+        for (LogicalPlan child : plan.children()) {
+            if (child instanceof SortAgnostic) {
+                OrderBy found = findOrderByNotPrecededByLimit(child);
+                if (found != null) {
+                    return found;
+                }
+            }
+        }
+        return null;
+    }
+
+    // Removes the found OrderBy node from its current position in the subtree
+    private static LogicalPlan removeOrderBy(LogicalPlan plan, OrderBy orderBy) {
+        if (plan == orderBy) {
+            return orderBy.child();
+        }
+        if (plan.children().isEmpty()) {
+            return plan;
+        }
+        return plan.replaceChildren(plan.children().stream().map(child -> removeOrderBy(child, orderBy)).toList());
+    }
+}

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Sample.java

@@ -18,7 +18,7 @@
 import java.io.IOException;
 import java.util.Objects;
 
-public class Sample extends UnaryPlan implements TelemetryAware {
+public class Sample extends UnaryPlan implements SortAgnostic, TelemetryAware {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(LogicalPlan.class, "Sample", Sample::new);
 
     private final Expression probability;