Add SortAware interface

luigidellaquila · luigidellaquila · commit 319cc69412df · 2025-02-06T11:19:49.000+01:00
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PruneRedundantOrderBy.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PruneRedundantOrderBy.java
@@ -8,25 +8,15 @@
 package org.elasticsearch.xpack.esql.optimizer.rules.logical;
 
 import org.elasticsearch.xpack.esql.plan.logical.Aggregate;
-import org.elasticsearch.xpack.esql.plan.logical.Drop;
-import org.elasticsearch.xpack.esql.plan.logical.Enrich;
-import org.elasticsearch.xpack.esql.plan.logical.Eval;
-import org.elasticsearch.xpack.esql.plan.logical.Filter;
-import org.elasticsearch.xpack.esql.plan.logical.InlineStats;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
-import org.elasticsearch.xpack.esql.plan.logical.Lookup;
-import org.elasticsearch.xpack.esql.plan.logical.MvExpand;
 import org.elasticsearch.xpack.esql.plan.logical.OrderBy;
-import org.elasticsearch.xpack.esql.plan.logical.Project;
-import org.elasticsearch.xpack.esql.plan.logical.RegexExtract;
-import org.elasticsearch.xpack.esql.plan.logical.Rename;
+import org.elasticsearch.xpack.esql.plan.logical.SortAware;
 import org.elasticsearch.xpack.esql.plan.logical.TopN;
 import org.elasticsearch.xpack.esql.plan.logical.UnaryPlan;
-import org.elasticsearch.xpack.esql.plan.logical.join.Join;
 
-import java.util.ArrayList;
+import java.util.ArrayDeque;
+import java.util.Deque;
 import java.util.IdentityHashMap;
-import java.util.List;
 
 /**
  * SORT cannot be executed without a LIMIT, as ES|QL doesn't support unbounded sort (yet).
@@ -56,7 +46,7 @@ protected LogicalPlan rule(LogicalPlan plan) {
             if (redundant.isEmpty()) {
                 return plan;
             }
-            return plan.transformUp(p -> {
+            return plan.transformDown(p -> {
                 if (redundant.containsKey(p)) {
                     return ((OrderBy) p).child();
                 }
@@ -67,50 +57,27 @@ protected LogicalPlan rule(LogicalPlan plan) {
         }
     }
 
+    /**
+     * breadth-first recursion to find redundant SORTs in the children tree.
+     */
     private IdentityHashMap<OrderBy, Void> findRedundantSort(LogicalPlan plan) {
-        List<LogicalPlan> toCheck = new ArrayList<>();
-        toCheck.add(plan);
-
         IdentityHashMap<OrderBy, Void> result = new IdentityHashMap<>();
-        LogicalPlan p = null;
+
+        Deque<LogicalPlan> toCheck = new ArrayDeque<>();
+        toCheck.push(plan);
+
         while (true) {
-            if (p == null) {
-                if (toCheck.isEmpty()) {
-                    return result;
-                } else {
-                    p = toCheck.remove(0);
-                }
-            } else if (p instanceof OrderBy ob) {
+            if (toCheck.isEmpty()) {
+                return result;
+            }
+            LogicalPlan p = toCheck.pop();
+            if (p instanceof OrderBy ob) {
                 result.put(ob, null);
-                p = ob.child();
-            } else if (p instanceof UnaryPlan unary) {
-                if (unary instanceof Project
-                    || unary instanceof Drop
-                    || unary instanceof Rename
-                    || unary instanceof MvExpand
-                    || unary instanceof Enrich
-                    || unary instanceof RegexExtract
-                    || unary instanceof InlineStats
-                    || unary instanceof Lookup
-                // IMPORTANT
-                // If we introduce window functions or order-sensitive aggs (eg. STREAMSTATS),
-                // the previous sort could actually become relevant
-                // so we have to be careful with plans that could use them, ie. the following
-                    || unary instanceof Filter
-                    || unary instanceof Eval
-                    || unary instanceof Aggregate) {
-                    p = unary.child();
-                } else {
-                    // stop here, other unary plans could be sensitive to SORT
-                    p = null;
+                toCheck.push(ob.child());
+            } else if (p instanceof SortAware sa && sa.dependsOnInputOrder() == false) {
+                for (LogicalPlan child : p.children()) {
+                    toCheck.push(child);
                 }
-            } else if (p instanceof Join lj) {
-                toCheck.add(lj.left());
-                toCheck.add(lj.right());
-                p = null;
-            } else {
-                // stop here, other unary plans could be sensitive to SORT
-                p = null;
             }
         }
     }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Aggregate.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Aggregate.java
@@ -40,7 +40,7 @@
 import static org.elasticsearch.xpack.esql.expression.NamedExpressions.mergeOutputAttributes;
 import static org.elasticsearch.xpack.esql.plan.logical.Filter.checkFilterConditionDataType;
 
-public class Aggregate extends UnaryPlan implements PostAnalysisVerificationAware, TelemetryAware {
+public class Aggregate extends UnaryPlan implements PostAnalysisVerificationAware, TelemetryAware, SortAware {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
         LogicalPlan.class,
         "Aggregate",
@@ -459,4 +459,10 @@ private static void addFailureOnGroupingUsedNakedInAggs(Failures failures, Expre
             fail(e, "grouping {} [{}] cannot be used as an aggregate once declared in the STATS BY clause", element, e.sourceText())
         );
     }
+
+    @Override
+    public boolean dependsOnInputOrder() {
+        // TODO review this if we introduce expressions that depend on input order (eg. window functions)
+        return SortAware.super.dependsOnInputOrder();
+    }
 }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Drop.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Drop.java
@@ -17,7 +17,7 @@
 import java.util.List;
 import java.util.Objects;
 
-public class Drop extends UnaryPlan implements TelemetryAware {
+public class Drop extends UnaryPlan implements TelemetryAware, SortAware {
     private final List<NamedExpression> removals;
 
     public Drop(Source source, LogicalPlan child, List<NamedExpression> removals) {
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Enrich.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Enrich.java
@@ -49,7 +49,7 @@
 import static org.elasticsearch.xpack.esql.core.expression.Expressions.asAttributes;
 import static org.elasticsearch.xpack.esql.expression.NamedExpressions.mergeOutputAttributes;
 
-public class Enrich extends UnaryPlan implements GeneratingPlan<Enrich>, PostAnalysisPlanVerificationAware, TelemetryAware {
+public class Enrich extends UnaryPlan implements GeneratingPlan<Enrich>, PostAnalysisPlanVerificationAware, TelemetryAware, SortAware {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
         LogicalPlan.class,
         "Enrich",
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Eval.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Eval.java
@@ -38,7 +38,7 @@
 import static org.elasticsearch.xpack.esql.core.expression.Expressions.asAttributes;
 import static org.elasticsearch.xpack.esql.expression.NamedExpressions.mergeOutputAttributes;
 
-public class Eval extends UnaryPlan implements GeneratingPlan<Eval>, PostAnalysisVerificationAware, TelemetryAware {
+public class Eval extends UnaryPlan implements GeneratingPlan<Eval>, PostAnalysisVerificationAware, TelemetryAware, SortAware {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(LogicalPlan.class, "Eval", Eval::new);
 
     private final List<Alias> fields;
@@ -189,4 +189,10 @@ public void postAnalysisVerification(Failures failures) {
             });
         });
     }
+
+    @Override
+    public boolean dependsOnInputOrder() {
+        // TODO review this if we introduce expressions that depend on input order (eg. window functions)
+        return SortAware.super.dependsOnInputOrder();
+    }
 }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Filter.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Filter.java
@@ -29,7 +29,7 @@
  * {@code SELECT x FROM y WHERE z ..} the "WHERE" clause is a Filter. A
  * {@code Filter} has a "condition" Expression that does the filtering.
  */
-public class Filter extends UnaryPlan implements PostAnalysisVerificationAware, TelemetryAware {
+public class Filter extends UnaryPlan implements PostAnalysisVerificationAware, TelemetryAware, SortAware {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(LogicalPlan.class, "Filter", Filter::new);
 
     private final Expression condition;
@@ -116,4 +116,10 @@ public static void checkFilterConditionDataType(Expression expression, Failures
             failures.add(fail(expression, "Condition expression needs to be boolean, found [{}]", expression.dataType()));
         }
     }
+
+    @Override
+    public boolean dependsOnInputOrder() {
+        // TODO review this if we introduce expressions that depend on input order (eg. window functions)
+        return SortAware.super.dependsOnInputOrder();
+    }
 }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/InlineStats.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/InlineStats.java
@@ -37,7 +37,7 @@
  *     underlying aggregate.
  * </p>
  */
-public class InlineStats extends UnaryPlan implements NamedWriteable, SurrogateLogicalPlan, TelemetryAware {
+public class InlineStats extends UnaryPlan implements NamedWriteable, SurrogateLogicalPlan, TelemetryAware, SortAware {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
         LogicalPlan.class,
         "InlineStats",
@@ -143,4 +143,10 @@ public boolean equals(Object obj) {
         InlineStats other = (InlineStats) obj;
         return Objects.equals(aggregate, other.aggregate);
     }
+
+    @Override
+    public boolean dependsOnInputOrder() {
+        // TODO review this if we introduce expressions that depend on input order (eg. window functions)
+        return SortAware.super.dependsOnInputOrder();
+    }
 }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Keep.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Keep.java
@@ -15,7 +15,7 @@
 import java.util.List;
 import java.util.Objects;
 
-public class Keep extends Project implements TelemetryAware {
+public class Keep extends Project implements TelemetryAware, SortAware {
 
     public Keep(Source source, LogicalPlan child, List<? extends NamedExpression> projections) {
         super(source, child, projections);
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Lookup.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Lookup.java
@@ -32,7 +32,7 @@
  * Looks up values from the associated {@code tables}.
  * The class is supposed to be substituted by a {@link Join}.
  */
-public class Lookup extends UnaryPlan implements SurrogateLogicalPlan, TelemetryAware {
+public class Lookup extends UnaryPlan implements SurrogateLogicalPlan, TelemetryAware, SortAware {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(LogicalPlan.class, "Lookup", Lookup::new);
 
     private final Expression tableName;
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/MvExpand.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/MvExpand.java
@@ -23,7 +23,7 @@
 import java.util.List;
 import java.util.Objects;
 
-public class MvExpand extends UnaryPlan implements TelemetryAware {
+public class MvExpand extends UnaryPlan implements TelemetryAware, SortAware {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(LogicalPlan.class, "MvExpand", MvExpand::new);
 
     private final NamedExpression target;
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/OrderBy.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/OrderBy.java
@@ -26,7 +26,12 @@
 
 import static org.elasticsearch.xpack.esql.common.Failure.fail;
 
-public class OrderBy extends UnaryPlan implements PostAnalysisVerificationAware, PostOptimizationVerificationAware, TelemetryAware {
+public class OrderBy extends UnaryPlan
+    implements
+        PostAnalysisVerificationAware,
+        PostOptimizationVerificationAware,
+        TelemetryAware,
+        SortAware {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(LogicalPlan.class, "OrderBy", OrderBy::new);
 
     private final List<Order> order;
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Project.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Project.java
@@ -25,7 +25,7 @@
 /**
  * A {@code Project} is a {@code Plan} with one child. In {@code SELECT x FROM y}, the "SELECT" statement is a Project.
  */
-public class Project extends UnaryPlan {
+public class Project extends UnaryPlan implements SortAware {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(LogicalPlan.class, "Project", Project::new);
 
     private final List<? extends NamedExpression> projections;
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/RegexExtract.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/RegexExtract.java
@@ -24,7 +24,7 @@
 import static org.elasticsearch.xpack.esql.common.Failure.fail;
 import static org.elasticsearch.xpack.esql.expression.NamedExpressions.mergeOutputAttributes;
 
-public abstract class RegexExtract extends UnaryPlan implements GeneratingPlan<RegexExtract>, PostAnalysisVerificationAware {
+public abstract class RegexExtract extends UnaryPlan implements GeneratingPlan<RegexExtract>, PostAnalysisVerificationAware, SortAware {
     protected final Expression input;
     protected final List<Attribute> extractedFields;
 
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Rename.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Rename.java
@@ -21,7 +21,7 @@
 import java.util.List;
 import java.util.Objects;
 
-public class Rename extends UnaryPlan implements TelemetryAware {
+public class Rename extends UnaryPlan implements TelemetryAware, SortAware {
 
     private final List<Alias> renamings;
 
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/SortAware.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/SortAware.java
@@ -0,0 +1,38 @@
+/*
+ * 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.plan.logical;
+
+/**
+ * This interface is intended to check redundancy of a previous SORT.
+ *
+ * Eg. if a MY_COMMAND that implements this interface is used after a sort, and if dependsOnInputOrder() = false,
+ * then we can assume that
+ * <p>
+ * <code>
+ * | SORT x | MY_COMMAND
+ * </code>
+ * <p>
+ * is equivalent to
+ * <p>
+ * <code>
+ * | MY_COMMAND
+ * </code>
+ * <p>
+ *
+ * In all the other cases, eg. if the command does not implement this interface, or if dependsOnInputOrder() = true
+ * then we assume that the SORT is still relevant and cannot be pruned.
+ */
+public interface SortAware {
+
+    /**
+     * Returns true if the logic of this command could depend on the order of the input.
+     */
+    default boolean dependsOnInputOrder() {
+        return false;
+    }
+}
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/join/Join.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/join/Join.java
@@ -21,6 +21,7 @@
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
 import org.elasticsearch.xpack.esql.plan.logical.BinaryPlan;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
+import org.elasticsearch.xpack.esql.plan.logical.SortAware;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -32,7 +33,7 @@
 import static org.elasticsearch.xpack.esql.expression.NamedExpressions.mergeOutputAttributes;
 import static org.elasticsearch.xpack.esql.plan.logical.join.JoinTypes.LEFT;
 
-public class Join extends BinaryPlan implements PostAnalysisVerificationAware {
+public class Join extends BinaryPlan implements PostAnalysisVerificationAware, SortAware {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(LogicalPlan.class, "Join", Join::new);
 
     private final JoinConfig config;
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizerTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizerTests.java
@@ -96,6 +96,7 @@
 import org.elasticsearch.xpack.esql.index.IndexResolution;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.LiteralsOnTheRight;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.OptimizerRules;
+import org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneRedundantOrderBy;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownAndCombineLimits;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownEnrich;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownEval;
@@ -7508,4 +7509,25 @@ public void testUnboundedSortExpandFilter() {
         assertThat(e.getMessage(), containsString("line 2:5: Unbounded sort not supported yet [SORT x] please add a limit"));
     }
 
+    public void testPruneRedundantOrderBy() {
+        var rule = new PruneRedundantOrderBy();
+
+        var query = """
+            row x = [1,2,3], y = 1
+            | sort x
+            | mv_expand x
+            | sort x
+            | mv_expand x
+            | sort y
+            """;
+        LogicalPlan analyzed = analyzer.analyze(parser.createStatement(query));
+        LogicalPlan optimized = rule.apply(analyzed);
+
+        // check that all the redundant SORTs are removed in a single run
+        var limit = as(optimized, Limit.class);
+        var orderBy = as(limit.child(), OrderBy.class);
+        var mvExpand = as(orderBy.child(), MvExpand.class);
+        var mvExpand2 = as(mvExpand.child(), MvExpand.class);
+        as(mvExpand2.child(), Row.class);
+    }
 }

Original file line number	Diff line number	Diff line change
`@@ -29,7 +29,7 @@`
`29`	`29`	`* {@code SELECT x FROM y WHERE z ..} the "WHERE" clause is a Filter. A`
`30`	`30`	`* {@code Filter} has a "condition" Expression that does the filtering.`
`31`	`31`	`*/`
`32`		`-public class Filter extends UnaryPlan implements PostAnalysisVerificationAware, TelemetryAware {`
	`32`	`+public class Filter extends UnaryPlan implements PostAnalysisVerificationAware, TelemetryAware, SortAware {`
`33`	`33`	`public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(LogicalPlan.class, "Filter", Filter::new);`
`34`	`34`
`35`	`35`	`private final Expression condition;`
`@@ -116,4 +116,10 @@ public static void checkFilterConditionDataType(Expression expression, Failures`
`116`	`116`	`failures.add(fail(expression, "Condition expression needs to be boolean, found [{}]", expression.dataType()));`
`117`	`117`	`}`
`118`	`118`	`}`
	`119`	`+`
	`120`	`+ @Override`
	`121`	`+ public boolean dependsOnInputOrder() {`
	`122`	`+ // TODO review this if we introduce expressions that depend on input order (eg. window functions)`
	`123`	`+ return SortAware.super.dependsOnInputOrder();`
	`124`	`+ }`
`119`	`125`	`}`