[ESQL] Adds planner implementation for LIMIT BY

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

@@ -133,6 +133,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;
@@ -230,6 +231,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;
@@ -818,6 +820,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,274 @@ emp_no:integer
 10004
 10005
 ;
+
+//
+// LIMIT BY
+//
+
+limitBy
+required_capability: esql_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: esql_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: esql_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: esql_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: esql_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: esql_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: esql_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: esql_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: esql_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: esql_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: esql_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
+;
+
+limitByWithLookupJoin
+required_capability: esql_limit_by
+required_capability: join_lookup_v12
+
+FROM employees
+| WHERE emp_no IN (10001, 10002, 10003, 10004, 10005)
+| EVAL language_code = languages
+| LOOKUP JOIN languages_lookup_non_unique_key ON language_code
+| SORT emp_no, language_code, language_name
+| LIMIT 1000
+| LIMIT 2 BY emp_no
+| KEEP emp_no, language_code, language_name
+;
+
+emp_no:integer | language_code:integer | language_name:keyword
+10001          | 2                     | German
+10001          | 2                     | German
+10002          | 5                     | null
+10003          | 4                     | Quenya
+10004          | 5                     | null
+10005          | 1                     | English
+10005          | 1                     | English
+;
+
+limitByShadowedNonJoinField
+required_capability: esql_limit_by
+required_capability: join_lookup_v12
+
+FROM employees
+| WHERE emp_no IN (10001, 10002, 10003, 10004, 10005)
+| EVAL language_code = languages
+| EVAL language_name = 2*salary
+| LOOKUP JOIN languages_lookup_non_unique_key ON language_code
+| SORT emp_no, language_code, language_name
+| LIMIT 1000
+| LIMIT 2 BY language_name
+| KEEP emp_no, language_code, language_name
+;
+
+emp_no:integer | language_code:integer | language_name:keyword
+10001          | 2                     | German
+10001          | 2                     | German
+10002          | 5                     | null
+10003          | 4                     | Quenya
+10004          | 5                     | null
+10005          | 1                     | English
+10005          | 1                     | English
+;
+
+limitByShadowedJoinField
+required_capability: esql_limit_by
+required_capability: join_lookup_v12
+
+FROM employees
+| WHERE emp_no IN (10001, 10002, 10003, 10004, 10005)
+| EVAL language_code = languages
+| LOOKUP JOIN languages_lookup_non_unique_key ON language_code
+| SORT emp_no, language_code, language_name
+| LIMIT 1000
+| LIMIT 2 BY language_code
+| KEEP emp_no, language_code, language_name
+;
+
+emp_no:integer | language_code:integer | language_name:keyword
+10001          | 2                     | German
+10001          | 2                     | German
+10002          | 5                     | null
+10003          | 4                     | Quenya
+10004          | 5                     | null
+10005          | 1                     | English
+10005          | 1                     | English
+;

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

@@ -225,7 +225,7 @@ limitCommand
     ;
 
 limitByGroupKey:
-    {this.isDevVersion()}? BY grouping=fields
+    {this.isDevVersion()}? BY booleanExpression (COMMA booleanExpression)*
     ;
 
 sortCommand

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

@@ -2321,8 +2321,10 @@ public enum Cap {
 
         /**
          * LIMIT n BY expr1, expr2 support for retaining at most n docs per group.
+         * Enables the feature without a preceding SORT.
+         *
          */
-        LIMIT_BY(Build.current().isSnapshot()),
+        ESQL_LIMIT_BY(Build.current().isSnapshot()),
 
         /**
          * Fix window validation in time-series aggregations when TBUCKET uses a numeric target bucket count.

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

@@ -40,8 +40,10 @@
 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.OrderBy;
 import org.elasticsearch.xpack.esql.plan.logical.Project;
 import org.elasticsearch.xpack.esql.plan.logical.Subquery;
+import org.elasticsearch.xpack.esql.plan.logical.UnaryPlan;
 import org.elasticsearch.xpack.esql.plan.logical.UnionAll;
 import org.elasticsearch.xpack.esql.plan.logical.promql.PromqlCommand;
 import org.elasticsearch.xpack.esql.telemetry.FeatureMetric;
@@ -357,9 +359,20 @@ 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 LimitBy) {
-            failures.add(fail(plan, "LIMIT BY is not yet supported"));
+        if (plan instanceof LimitBy limitBy) {
+            LogicalPlan child = limitBy.child();
+            while (child instanceof UnaryPlan unary) {
+                if (child instanceof OrderBy) {
+                    failures.add(fail(limitBy, "SORT cannot be used before LIMIT BY"));
+                    break;
+                }
+                if (child instanceof Limit) {
+                    break;
+                }
+                child = unary.child();
+            }
         }
     }
 

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

@@ -188,6 +188,23 @@ 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);
+            if (left == null || right == null) {
+                throw new IllegalArgumentException("Unexpected null expression in list at index [" + i + "]");
+            }
+            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<>();