ESQL: Adds GroupedTopN command

docs/changelog/140019.yaml

@@ -0,0 +1,5 @@
+pr: 140019
+summary: Adds `GroupedTopN` ESQL parsing and planning
+area: ES|QL
+type: feature
+issues: []

server/src/main/resources/transport/definitions/referable/esql_limit_by.csv

@@ -0,0 +1 @@
+9314000

server/src/main/resources/transport/upper_bounds/9.4.csv

@@ -1 +1 @@
-sql_optional_allow_partial_search_results,9313000
+esql_limit_by,9314000

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java

@@ -2205,6 +2205,12 @@ public enum Cap {
          */
         EXTERNAL_COMMAND(Build.current().isSnapshot()),
 
+        /**
+
+         * Enables LIMIT N BY in the LIMIT command, both with and without a preceding SORT.
+         */
+        LIMIT_BY,
+
         /**
          * https://github.com/elastic/elasticsearch/issues/142219
          */

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Analyzer.java

@@ -1911,7 +1911,8 @@ private InferenceFunction<?> resolveInferenceFunction(InferenceFunction<?> infer
     private static class AddImplicitLimit extends ParameterizedRule<LogicalPlan, LogicalPlan, AnalyzerContext> {
         @Override
         public LogicalPlan apply(LogicalPlan logicalPlan, AnalyzerContext context) {
-            List<LogicalPlan> limits = logicalPlan.collectFirstChildren(Limit.class::isInstance);
+            // Find existing LIMITs without groups. Groups are unbounded, and they still require another default LIMIT
+            List<LogicalPlan> limits = logicalPlan.collectFirstChildren(lp -> lp instanceof Limit l && l.groupings().isEmpty());
             // We check whether the query contains a TimeSeriesAggregate to determine if we should apply
             // the default limit for TS queries or for non-TS queries.
             // NOTE: PromqlCommand is translated to TimeSeriesAggregate during optimization.

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Verifier.java

@@ -34,7 +34,6 @@
 import org.elasticsearch.xpack.esql.plan.logical.InlineStats;
 import org.elasticsearch.xpack.esql.plan.logical.Insist;
 import org.elasticsearch.xpack.esql.plan.logical.Limit;
-import org.elasticsearch.xpack.esql.plan.logical.LimitBy;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.esql.plan.logical.Lookup;
 import org.elasticsearch.xpack.esql.plan.logical.Project;
@@ -120,7 +119,6 @@ Collection<Failure> verify(LogicalPlan plan, BitSet partialMetrics) {
             checkUnsupportedAttributeRenaming(p, failures);
             checkInsist(p, failures);
             checkLimitBeforeInlineStats(p, failures);
-            checkLimitBy(p, failures);
         });
 
         if (failures.hasFailures() == false) {
@@ -334,12 +332,6 @@ private static void checkLimitBeforeInlineStats(LogicalPlan plan, Failures failu
         }
     }
 
-    private static void checkLimitBy(LogicalPlan plan, Failures failures) {
-        if (plan instanceof LimitBy) {
-            failures.add(fail(plan, "LIMIT BY is not yet supported"));
-        }
-    }
-
     private void licenseCheck(LogicalPlan plan, Failures failures) {
         Consumer<Node<?>> licenseCheck = n -> {
             if (n instanceof LicenseAware la && la.licenseCheck(licenseState) == false) {

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/core/expression/Expressions.java

@@ -188,6 +188,21 @@ public static boolean equalsAsAttribute(Expression left, Expression right) {
         return true;
     }
 
+    public static boolean listSemanticEquals(List<Expression> leftList, List<Expression> rightList) {
+        if (leftList.size() != rightList.size()) {
+            return false;
+        }
+        for (int i = 0; i < leftList.size(); i++) {
+            Expression left = leftList.get(i);
+            Expression right = rightList.get(i);
+            assert left != null && right != null;
+            if (left.semanticEquals(right) == false) {
+                return false;
+            }
+        }
+        return true;
+    }
+
     public static List<Tuple<Attribute, Expression>> aliases(List<? extends NamedExpression> named) {
         // an alias of same name and data type can be reused (by mistake): need to use a list to collect all refs (and later report them)
         List<Tuple<Attribute, Expression>> aliases = new ArrayList<>();

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

@@ -60,6 +60,7 @@
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.ReplaceAggregateNestedExpressionWithEval;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.ReplaceAliasingEvalWithProject;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.ReplaceLimitAndSortAsTopN;
+import org.elasticsearch.xpack.esql.optimizer.rules.logical.ReplaceLimitByExpressionWithEval;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.ReplaceOrderByExpressionWithEval;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.ReplaceRegexMatch;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.ReplaceRowAsLocalRelation;
@@ -181,6 +182,7 @@ protected static Batch<LogicalPlan> substitutions() {
             // check for a trivial conversion introduced by a surrogate
             new ReplaceTrivialTypeConversions(),
             new ReplaceOrderByExpressionWithEval(),
+            new ReplaceLimitByExpressionWithEval(),
             // new NormalizeAggregate(), - waits on https://github.com/elastic/elasticsearch/issues/100634
             new SubstituteApproximationPlan()
         );

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

@@ -13,6 +13,8 @@
 import org.elasticsearch.xpack.esql.plan.logical.Project;
 import org.elasticsearch.xpack.esql.plan.logical.TopN;
 
+import static org.elasticsearch.xpack.esql.core.expression.Expressions.listSemanticEquals;
+
 /**
  * Combines a Limit immediately followed by a TopN into a single TopN.
  * This is needed because {@link HoistRemoteEnrichTopN} can create new TopN nodes that are not covered by the previous rules.
@@ -25,13 +27,13 @@ public CombineLimitTopN() {
 
     @Override
     public LogicalPlan rule(Limit limit) {
-        if (limit.child() instanceof TopN topn) {
+        if (limit.child() instanceof TopN topn && listSemanticEquals(limit.groupings(), topn.groupings())) {
             int thisLimitValue = Foldables.limitValue(limit.limit(), limit.sourceText());
             int topNValue = Foldables.limitValue(topn.limit(), topn.sourceText());
             if (topNValue <= thisLimitValue) {
                 return topn;
             } else {
-                return new TopN(topn.source(), topn.child(), topn.order(), limit.limit(), topn.local());
+                return new TopN(topn.source(), topn.child(), topn.order(), limit.limit(), topn.groupings(), topn.local());
             }
         }
         if (limit.child() instanceof Project proj) {

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

@@ -60,7 +60,7 @@ protected LogicalPlan rule(Enrich en) {
                     if (missingAttributes.isEmpty()) {
                         // No need for aliasing - then it's simple, just copy the TopN on top of Enrich and mark the original as local
                         LogicalPlan transformedEnrich = en.transformDown(TopN.class, t -> t == topN ? topN.withLocal(true) : t);
-                        return new TopN(topN.source(), transformedEnrich, topN.order(), topN.limit(), false);
+                        return new TopN(topN.source(), transformedEnrich, topN.order(), topN.limit(), topN.groupings(), false);
                     } else {
                         Set<String> missingNames = new HashSet<>(Expressions.names(missingAttributes));
                         AttributeMap.Builder<Alias> aliasesForReplacedAttributesBuilder = AttributeMap.builder();
@@ -127,7 +127,7 @@ protected LogicalPlan rule(Enrich en) {
                         }
 
                         // Create a copy of TopN with the rewritten attributes on top on Enrich
-                        var copyTop = new TopN(topN.source(), transformedEnrich, newOrder, topN.limit(), false);
+                        var copyTop = new TopN(topN.source(), transformedEnrich, newOrder, topN.limit(), topN.groupings(), false);
                         // And use the Project to remove the fields that we don't need anymore
                         return new Project(en.source(), copyTop, outputs);
                     }

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

@@ -39,9 +39,10 @@ public LogicalPlan apply(LogicalPlan plan, LogicalOptimizerContext ctx) {
         AttributeMap.Builder<Expression> builder = AttributeMap.builder();
         builder.putAll(collectRefs);
 
-        // Exclude multi-value (List) literals used as GROUP BY keys from propagation.
+        // Exclude multi-value (List) literals used as grouping keys from propagation.
         //
-        // Rationale: GROUP BY explodes multi-value fields into single values. For example:
+        // Rationale: both GROUP BY and LIMIT BY explode multi-value fields into single values.
+        // For example:
         // ROW a = [1, 2] | STATS x = a + SUM(a) BY a
         // Before aggregation, `a` is multi-valued [1, 2]. After GROUP BY, `a` becomes single-valued
         // (either 1 or 2 per group). If we propagate the original multi-value literal [1, 2] into

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

@@ -69,7 +69,7 @@ private static LogicalPlan pruneColumns(LogicalPlan plan, AttributeSet.Builder u
 
             // TODO: revisit with every new command
             // skip nodes that simply pass the input through and use no references
-            if (p instanceof Limit || p instanceof Sample) {
+            if (p instanceof Limit limit && limit.groupings().isEmpty() || p instanceof Sample) {
                 return p;
             }
 

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

@@ -8,7 +8,10 @@
 package org.elasticsearch.xpack.esql.optimizer.rules.logical;
 
 import org.elasticsearch.xpack.esql.core.expression.Alias;
+import org.elasticsearch.xpack.esql.core.expression.Attribute;
+import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.FoldContext;
+import org.elasticsearch.xpack.esql.core.expression.NameId;
 import org.elasticsearch.xpack.esql.core.util.Holder;
 import org.elasticsearch.xpack.esql.expression.function.fulltext.Score;
 import org.elasticsearch.xpack.esql.optimizer.LogicalOptimizerContext;
@@ -31,7 +34,11 @@
 import org.elasticsearch.xpack.esql.rule.Rule;
 
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Set;
+
+import static org.elasticsearch.xpack.esql.core.expression.Expressions.listSemanticEquals;
 
 public final class PushDownAndCombineLimits extends OptimizerRules.ParameterizedOptimizerRule<Limit, LogicalOptimizerContext>
     implements
@@ -50,7 +57,7 @@ private PushDownAndCombineLimits(boolean local) {
 
     @Override
     public LogicalPlan rule(Limit limit, LogicalOptimizerContext ctx) {
-        if (limit.child() instanceof Limit childLimit) {
+        if (limit.child() instanceof Limit childLimit && listSemanticEquals(childLimit.groupings(), limit.groupings())) {
             return combineLimits(limit, childLimit, ctx.foldCtx());
         } else if (limit.child() instanceof UnaryPlan unary) {
             if (unary instanceof Eval
@@ -61,6 +68,8 @@ public LogicalPlan rule(Limit limit, LogicalOptimizerContext ctx) {
                 if (false == local && unary instanceof Eval && evalAliasNeedsData((Eval) unary)) {
                     // do not push down the limit through an eval that needs data (e.g. a score function) during initial planning
                     return limit;
+                } else if (groupingsReferenceAttributeDefinedByChild(limit, unary)) {
+                    return limit;
                 } else {
                     return unary.replaceChild(limit.replaceChild(unary.child()));
                 }
@@ -82,7 +91,7 @@ public LogicalPlan rule(Limit limit, LogicalOptimizerContext ctx) {
             // this applies for cases such as | limit 1 | sort field | limit 10
             else {
                 Limit descendantLimit = descendantLimit(unary);
-                if (descendantLimit != null) {
+                if (descendantLimit != null && descendantLimit.groupings().isEmpty()) {
                     var l1 = (int) limit.limit().fold(ctx.foldCtx());
                     var l2 = (int) descendantLimit.limit().fold(ctx.foldCtx());
                     if (l2 <= l1) {
@@ -163,7 +172,7 @@ private static Limit combineLimits(Limit upper, Limit lower, FoldContext ctx) {
             // If any of them is local, we want the local limit
             return lower.local() ? lower : lower.withLocal(upper.local());
         } else {
-            return new Limit(upper.source(), upper.limit(), lower.child(), upper.duplicated(), upper.local());
+            return new Limit(upper.source(), upper.limit(), lower.child(), upper.groupings(), upper.duplicated(), upper.local());
         }
     }
 
@@ -183,14 +192,35 @@ private boolean evalAliasNeedsData(Eval eval) {
     }
 
     /**
-     * Checks the existence of another 'visible' Limit, that exists behind an operation that doesn't produce output more data than
-     * its input (that is not a relation/source nor aggregation).
+     * Returns {@code true} if any attribute referenced by the limit's groupings is defined by the child node
+     * (i.e. present in the child's output but absent from the grandchild's output). Pushing a grouped limit
+     * past such a child would leave the grouping attribute unresolved.
+     */
+    private static boolean groupingsReferenceAttributeDefinedByChild(Limit limit, UnaryPlan child) {
+        if (limit.groupings().isEmpty()) {
+            return false;
+        }
+        Set<NameId> grandchildIds = new HashSet<>();
+        for (Attribute a : child.child().output()) {
+            grandchildIds.add(a.id());
+        }
+        for (Expression g : limit.groupings()) {
+            if (g instanceof Attribute a && grandchildIds.contains(a.id()) == false) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    /**
+     * Checks the existence of another 'visible' non-grouping Limit, that exists behind an operation that doesn't produce more output data
+     * than its input (that is not a relation/source nor aggregation).
      * P.S. Typically an aggregation produces less data than the input.
      */
     private static Limit descendantLimit(UnaryPlan unary) {
         UnaryPlan plan = unary;
         while (plan instanceof Aggregate == false) {
-            if (plan instanceof Limit limit) {
+            if (plan instanceof Limit limit && limit.groupings().isEmpty()) {
                 return limit;
             } else if (plan instanceof MvExpand) {
                 // the limit that applies to mv_expand shouldn't be changed

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

@@ -15,11 +15,13 @@
 public final class ReplaceLimitAndSortAsTopN extends OptimizerRules.OptimizerRule<Limit> {
 
     @Override
-    protected LogicalPlan rule(Limit plan) {
-        LogicalPlan p = plan;
-        if (plan.child() instanceof OrderBy o) {
-            p = new TopN(o.source(), o.child(), o.order(), plan.limit(), false);
+    protected LogicalPlan rule(Limit limit) {
+        LogicalPlan p = limit;
+
+        if (limit.child() instanceof OrderBy o) {
+            p = new TopN(o.source(), o.child(), o.order(), limit.limit(), limit.groupings(), false);
         }
+
         return p;
     }
 }