[ESQL] Adds parsing and planner wiring for LIMIT BY

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/build.gradle

@@ -431,7 +431,6 @@ tasks.register("regenParser", JavaExec) {
     '-package', 'org.elasticsearch.xpack.esql.parser',
     '-listener',
     '-visitor',
-    '-lib', outputPath,
     '-lib', "${file(grammarPath)}/parser",
     '-o', outputPath,
     "${file(grammarPath)}/EsqlBaseParser.g4"

x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/EsqlTestUtils.java

@@ -131,6 +131,7 @@
 import org.elasticsearch.xpack.esql.plan.QuerySettings;
 import org.elasticsearch.xpack.esql.plan.logical.Enrich;
 import org.elasticsearch.xpack.esql.plan.logical.EsRelation;
+import org.elasticsearch.xpack.esql.plan.logical.Eval;
 import org.elasticsearch.xpack.esql.plan.logical.Explain;
 import org.elasticsearch.xpack.esql.plan.logical.Limit;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
@@ -227,6 +228,7 @@
 import static org.elasticsearch.xpack.esql.parser.ParserUtils.ParamClassification.VALUE;
 import static org.elasticsearch.xpack.esql.plan.QuerySettings.UNMAPPED_FIELDS;
 import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.hasSize;
 import static org.hamcrest.Matchers.instanceOf;
 import static org.junit.Assert.assertNotNull;
@@ -758,6 +760,20 @@ public static Limit asLimit(Object node, Integer limitLiteral, Boolean duplicate
         return limit;
     }
 
+    /**
+     * Assert that an {@link Eval}'s fields are literal-valued aliases with the given names and values (in order).
+     */
+    public static Eval assertEvalFields(Eval eval, String[] names, Object[] values) {
+        var fields = eval.fields();
+        Assert.assertEquals(names.length, fields.size());
+        Assert.assertEquals(names.length, values.length);
+        for (int i = 0; i < names.length; i++) {
+            assertThat(fields.get(i).name(), equalTo(names[i]));
+            assertThat(as(fields.get(i).child(), Literal.class).value(), equalTo(values[i]));
+        }
+        return eval;
+    }
+
     public static Map<String, EsField> loadMapping(String name) {
         return LoadMapping.loadMapping(name);
     }

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

@@ -14,3 +14,201 @@ emp_no:integer
 10004
 10005
 ;
+
+//
+// LIMIT BY
+//
+
+limitBy
+required_capability: limit_by
+
+FROM employees
+| WHERE emp_no IN (10001, 10002, 10003, 10005, 10006)
+| SORT emp_no
+| LIMIT 1000
+| LIMIT 5 BY languages
+| KEEP emp_no, first_name, languages
+;
+
+emp_no:integer | first_name:keyword | languages:integer
+10001          | Georgi             | 2
+10002          | Bezalel            | 5
+10003          | Parto              | 4
+10005          | Kyoichi            | 1
+10006          | Anneke             | 3
+;
+
+limitByLimit0
+required_capability: limit_by
+
+FROM employees
+| WHERE emp_no IN (10001, 10002, 10003)
+| SORT emp_no
+| LIMIT 1000
+| LIMIT 0 BY languages
+| KEEP emp_no, languages
+;
+
+emp_no:integer | languages:integer
+;
+
+limitByMultipleColumns
+required_capability: limit_by
+
+FROM employees
+| WHERE emp_no IN (10001, 10003, 10004, 10007)
+| SORT emp_no
+| LIMIT 1000
+| LIMIT 5 BY languages, gender
+| KEEP emp_no, first_name, languages, gender
+;
+
+emp_no:integer | first_name:keyword | languages:integer | gender:keyword
+10001          | Georgi             | 2                 | M
+10003          | Parto              | 4                 | M
+10004          | Chirstian          | 5                 | M
+10007          | Tzvetan            | 4                 | F
+;
+
+limitByWithExpression
+required_capability: limit_by
+
+FROM employees
+| WHERE emp_no IN (10001, 10002, 10003, 10004, 10005)
+| SORT emp_no
+| LIMIT 1000
+| LIMIT 5 BY languages * 2
+| KEEP emp_no, first_name, languages
+;
+
+emp_no:integer | first_name:keyword | languages:integer
+10001          | Georgi             | 2
+10002          | Bezalel            | 5
+10003          | Parto              | 4
+10004          | Chirstian          | 5
+10005          | Kyoichi            | 1
+;
+
+limitByMultivalueGroupKey
+required_capability: limit_by
+
+FROM employees
+| WHERE emp_no IN (10001, 10008)
+| SORT emp_no
+| LIMIT 1000
+| LIMIT 5 BY job_positions
+| KEEP emp_no, first_name, job_positions
+;
+
+emp_no:integer | first_name:keyword | job_positions:keyword
+10001          | Georgi             | [Accountant, Senior Python Developer]
+10008          | Saniya             | [Internship, Junior Developer, Purchase Manager, Senior Python Developer]
+;
+
+limitByNullGroup
+required_capability: limit_by
+
+FROM employees
+| WHERE emp_no IN (10001, 10020)
+| SORT emp_no
+| LIMIT 1000
+| LIMIT 5 BY languages
+| KEEP emp_no, first_name, languages
+;
+
+emp_no:integer | first_name:keyword | languages:integer
+10001          | Georgi             | 2
+10020          | Mayuko             | null
+;
+
+limitByConstant
+required_capability: limit_by
+
+FROM employees
+| WHERE emp_no IN (10001, 10002, 10003, 10005, 10006)
+| SORT emp_no
+| LIMIT 1000
+| LIMIT 2 BY 5*42
+| KEEP emp_no, first_name, languages
+;
+
+emp_no:integer | first_name:keyword | languages:integer
+10001          | Georgi             | 2
+10002          | Bezalel            | 5
+;
+
+limitByWithExpressionAndConstant
+required_capability: limit_by
+
+FROM employees
+| WHERE emp_no IN (10001, 10002, 10003, 10004, 10005)
+| SORT emp_no
+| LIMIT 1000
+| LIMIT 5 BY languages * 2, 20 - 5 * 2
+| KEEP emp_no, first_name, languages
+;
+
+emp_no:integer | first_name:keyword | languages:integer
+10001          | Georgi             | 2
+10002          | Bezalel            | 5
+10003          | Parto              | 4
+10004          | Chirstian          | 5
+10005          | Kyoichi            | 1
+;
+
+limitByWithAlias
+required_capability: limit_by
+
+FROM employees
+| WHERE emp_no IN (10001, 10002, 10003, 10004, 10005)
+| SORT emp_no
+| LIMIT 1000
+| EVAL g = languages * 2
+| LIMIT 5 BY g
+| KEEP emp_no, first_name, languages
+;
+
+emp_no:integer | first_name:keyword | languages:integer
+10001          | Georgi             | 2
+10002          | Bezalel            | 5
+10003          | Parto              | 4
+10004          | Chirstian          | 5
+10005          | Kyoichi            | 1
+;
+
+limitByThenStats
+required_capability: limit_by
+
+FROM employees
+| LIMIT 2 BY languages
+| STATS c = COUNT(*) BY languages
+| SORT languages ASC NULLS LAST
+;
+
+c:long | languages:integer
+2      | 1
+2      | 2
+2      | 3
+2      | 4
+2      | 5
+2      | null
+;
+
+limitByPrecededByStats
+required_capability: limit_by
+
+FROM employees
+| STATS cnt=COUNT(*) BY job_positions, languages
+| SORT job_positions, cnt DESC, languages
+| LIMIT 1000
+| LIMIT 1 BY job_positions
+| LIMIT 5
+;
+
+cnt:long | job_positions:keyword | languages:integer
+5        | Accountant            | 2
+5        | Architect             | 4
+4        | Business Analyst      | 2
+3        | Data Scientist        | 1
+5        | Head Human Resources  | 5
+;

x-pack/plugin/esql/src/main/antlr/EsqlBaseParser.g4

@@ -221,7 +221,11 @@ stringOrParameter
     ;
 
 limitCommand
-    : LIMIT constant
+    : LIMIT constant limitByGroupKey?
+    ;
+
+limitByGroupKey:
+    {this.isDevVersion()}? BY grouping=fields
     ;
 
 sortCommand

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

@@ -2190,6 +2190,12 @@ public enum Cap {
          */
         EXTERNAL_COMMAND(Build.current().isSnapshot()),
 
+        /**
+
+         * Enables LIMIT N BY 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/Verifier.java

@@ -36,7 +36,9 @@
 import org.elasticsearch.xpack.esql.plan.logical.Limit;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.esql.plan.logical.Lookup;
+import org.elasticsearch.xpack.esql.plan.logical.OrderBy;
 import org.elasticsearch.xpack.esql.plan.logical.Project;
+import org.elasticsearch.xpack.esql.plan.logical.UnaryPlan;
 import org.elasticsearch.xpack.esql.plan.logical.promql.PromqlCommand;
 import org.elasticsearch.xpack.esql.telemetry.FeatureMetric;
 import org.elasticsearch.xpack.esql.telemetry.Metrics;
@@ -119,6 +121,7 @@ Collection<Failure> verify(LogicalPlan plan, BitSet partialMetrics) {
             checkUnsupportedAttributeRenaming(p, failures);
             checkInsist(p, failures);
             checkLimitBeforeInlineStats(p, failures);
+            checkLimitBy(p, failures);
         });
 
         if (failures.hasFailures() == false) {
@@ -332,6 +335,23 @@ private static void checkLimitBeforeInlineStats(LogicalPlan plan, Failures failu
         }
     }
 
+    // TODO: remove this check when SORT + LIMIT BY (TopN) support is added
+    private static void checkLimitBy(LogicalPlan plan, Failures failures) {
+        if (plan instanceof Limit limit && limit.groupings().isEmpty() == false) {
+            LogicalPlan child = limit.child();
+            while (child instanceof UnaryPlan unary) {
+                if (child instanceof OrderBy) {
+                    failures.add(fail(limit, "SORT cannot be used before LIMIT BY"));
+                    break;
+                }
+                if (child instanceof Limit l && l.groupings().isEmpty()) {
+                    break;
+                }
+                child = unary.child();
+            }
+        }
+    }
+
     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

@@ -178,6 +178,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

@@ -37,6 +37,7 @@
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneEmptyAggregates;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneEmptyForkBranches;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneFilters;
+import org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneLiteralsInLimitBy;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneLiteralsInOrderBy;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneRedundantOrderBy;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneRedundantSortClauses;
@@ -60,6 +61,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 +183,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()
         );
@@ -224,6 +227,7 @@ protected static Batch<LogicalPlan> operators() {
             new PruneFilters(),
             new PruneColumns(),
             new PruneLiteralsInOrderBy(),
+            new PruneLiteralsInLimitBy(),
             new PushDownAndCombineLimits(),
             new PushLimitToKnn(),
             new PushDownAndCombineFilters(),

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

@@ -25,7 +25,8 @@ public CombineLimitTopN() {
 
     @Override
     public LogicalPlan rule(Limit limit) {
-        if (limit.child() instanceof TopN topn) {
+        // TODO: update this check when SORT + LIMIT BY (TopN) support is added
+        if (limit.groupings().isEmpty() && limit.child() instanceof TopN topn) {
             int thisLimitValue = Foldables.limitValue(limit.limit(), limit.sourceText());
             int topNValue = Foldables.limitValue(topn.limit(), topn.sourceText());
             if (topNValue <= thisLimitValue) {