[ESQL] Plug TopN agg groupings filtering into the language

benchmarks/src/main/java/org/elasticsearch/benchmark/compute/operator/AggregatorBenchmark.java

@@ -197,7 +197,7 @@ private static Operator operator(DriverContext driverContext, String grouping, S
         };
         return new HashAggregationOperator(
             List.of(supplier(op, dataType, filter).groupingAggregatorFactory(AggregatorMode.SINGLE, List.of(groups.size()))),
-            () -> BlockHash.build(groups, driverContext.blockFactory(), 16 * 1024, false),
+            () -> BlockHash.build(groups, driverContext.blockFactory(), 16 * 1024, false, TOP_N_LIMIT),
             driverContext
         );
     }

benchmarks/src/main/java/org/elasticsearch/benchmark/compute/operator/ValuesAggregatorBenchmark.java

@@ -123,7 +123,7 @@ private static Operator operator(DriverContext driverContext, int groups, String
         List<BlockHash.GroupSpec> groupSpec = List.of(new BlockHash.GroupSpec(0, ElementType.LONG));
         return new HashAggregationOperator(
             List.of(supplier(dataType).groupingAggregatorFactory(AggregatorMode.SINGLE, List.of(1))),
-            () -> BlockHash.build(groupSpec, driverContext.blockFactory(), 16 * 1024, false),
+            () -> BlockHash.build(groupSpec, driverContext.blockFactory(), 16 * 1024, false, 100),
             driverContext
         ) {
             @Override

docs/changelog/130111.yaml

@@ -0,0 +1,5 @@
+pr: 130111
+summary: Plug TopN agg groupings filtering into the language
+area: ES|QL
+type: enhancement
+issues: []

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/blockhash/BlockHash.java

@@ -159,12 +159,20 @@ public boolean isCategorize() {
      *                                 null handling and remove this flag, but we need to disable these in
      *                                 production until we can. And this lets us continue to compile and
      *                                 test them.
+     * @param maxTopNLimit the maximum limit for TopN groups to use a TopNBlockHash.
+     *                     This usually comes from {@code QueryPragma.maxTopNAggsLimit()}.
      */
-    public static BlockHash build(List<GroupSpec> groups, BlockFactory blockFactory, int emitBatchSize, boolean allowBrokenOptimizations) {
+    public static BlockHash build(
+        List<GroupSpec> groups,
+        BlockFactory blockFactory,
+        int emitBatchSize,
+        boolean allowBrokenOptimizations,
+        int maxTopNLimit
+    ) {
         if (groups.size() == 1) {
             GroupSpec group = groups.get(0);
-            if (group.topNDef() != null && group.elementType() == ElementType.LONG) {
-                TopNDef topNDef = group.topNDef();
+            TopNDef topNDef = group.topNDef();
+            if (topNDef != null && maxTopNLimit > 0 && group.elementType() == ElementType.LONG && topNDef.limit() < maxTopNLimit) {
                 return new LongTopNBlockHash(group.channel(), topNDef.asc(), topNDef.nullsFirst(), topNDef.limit(), blockFactory);
             }
             return newForElementType(group.channel(), group.elementType(), blockFactory);

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/HashAggregationOperator.java

@@ -46,7 +46,8 @@ public record HashAggregationOperatorFactory(
         AggregatorMode aggregatorMode,
         List<GroupingAggregator.Factory> aggregators,
         int maxPageSize,
-        AnalysisRegistry analysisRegistry
+        AnalysisRegistry analysisRegistry,
+        int maxTopNLimit
     ) implements OperatorFactory {
         @Override
         public Operator get(DriverContext driverContext) {
@@ -65,7 +66,7 @@ public Operator get(DriverContext driverContext) {
             }
             return new HashAggregationOperator(
                 aggregators,
-                () -> BlockHash.build(groups, driverContext.blockFactory(), maxPageSize, false),
+                () -> BlockHash.build(groups, driverContext.blockFactory(), maxPageSize, false, maxTopNLimit),
                 driverContext
             );
         }

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/TimeSeriesAggregationOperator.java

@@ -35,7 +35,8 @@ public record Factory(
         List<BlockHash.GroupSpec> groups,
         AggregatorMode aggregatorMode,
         List<GroupingAggregator.Factory> aggregators,
-        int maxPageSize
+        int maxPageSize,
+        int maxTopNLimit
     ) implements OperatorFactory {
         @Override
         public Operator get(DriverContext driverContext) {
@@ -48,7 +49,8 @@ public Operator get(DriverContext driverContext) {
                         groups,
                         driverContext.blockFactory(),
                         maxPageSize,
-                        true // we can enable optimizations as the inputs are vectors
+                        true, // we can enable optimizations as the inputs are vectors
+                        maxTopNLimit
                     );
                 }
             }, driverContext);

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/GroupingAggregatorFunctionTestCase.java

@@ -117,15 +117,17 @@ private Operator.OperatorFactory simpleWithMode(
                 mode,
                 List.of(supplier.groupingAggregatorFactory(mode, channels(mode))),
                 randomPageSize(),
-                null
+                null,
+                100
             );
         } else {
             return new RandomizingHashAggregationOperatorFactory(
                 List.of(new BlockHash.GroupSpec(0, ElementType.LONG)),
                 mode,
                 List.of(supplier.groupingAggregatorFactory(mode, channels(mode))),
                 randomPageSize(),
-                null
+                null,
+                100
             );
         }
     }
@@ -824,7 +826,8 @@ private record RandomizingHashAggregationOperatorFactory(
         AggregatorMode aggregatorMode,
         List<GroupingAggregator.Factory> aggregators,
         int maxPageSize,
-        AnalysisRegistry analysisRegistry
+        AnalysisRegistry analysisRegistry,
+        int maxTopNLimit
     ) implements Operator.OperatorFactory {
 
         @Override
@@ -838,7 +841,7 @@ public Operator get(DriverContext driverContext) {
                         analysisRegistry,
                         maxPageSize
                     )
-                    : BlockHash.build(groups, driverContext.blockFactory(), maxPageSize, false);
+                    : BlockHash.build(groups, driverContext.blockFactory(), maxPageSize, false, maxTopNLimit);
 
                 return new BlockHashWrapper(driverContext.blockFactory(), blockHash) {
                     @Override
@@ -886,7 +889,8 @@ public String describe() {
                 aggregatorMode,
                 aggregators,
                 maxPageSize,
-                analysisRegistry
+                analysisRegistry,
+                maxTopNLimit
             ).describe();
         }
     }

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/blockhash/BlockHashRandomizedTests.java

@@ -256,7 +256,7 @@ private BlockHash newBlockHash(BlockFactory blockFactory, int emitBatchSize, Lis
         }
         return forcePackedHash
             ? new PackedValuesBlockHash(specs, blockFactory, emitBatchSize)
-            : BlockHash.build(specs, blockFactory, emitBatchSize, true);
+            : BlockHash.build(specs, blockFactory, emitBatchSize, true, 100);
     }
 
     private static final int LOOKUP_POSITIONS = 1_000;

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/blockhash/BlockHashTests.java

@@ -47,6 +47,8 @@
 import static org.hamcrest.Matchers.endsWith;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThan;
+import static org.hamcrest.Matchers.instanceOf;
+import static org.hamcrest.Matchers.not;
 import static org.hamcrest.Matchers.startsWith;
 
 public class BlockHashTests extends BlockHashTestCase {
@@ -1544,7 +1546,8 @@ public void testTimeSeriesBlockHash() throws Exception {
             List.of(new BlockHash.GroupSpec(0, ElementType.BYTES_REF), new BlockHash.GroupSpec(1, ElementType.LONG)),
             blockFactory,
             32 * 1024,
-            forcePackedHash
+            forcePackedHash,
+            100
         );
         int numPages = between(1, 100);
         int globalTsid = -1;
@@ -1660,6 +1663,34 @@ public void close() {
         }
     }
 
+    public void testTopNBlockHashLimit() {
+        int limit = randomIntBetween(1, 100);
+
+        try (
+            var hash = BlockHash.build(
+                List.of(new BlockHash.GroupSpec(0, ElementType.LONG, null, new BlockHash.TopNDef(0, true, false, limit))),
+                blockFactory,
+                32 * 1024,
+                forcePackedHash,
+                randomIntBetween(limit, 1000)
+            )
+        ) {
+            assertThat(hash, instanceOf(LongTopNBlockHash.class));
+        }
+
+        try (
+            var hash = BlockHash.build(
+                List.of(new BlockHash.GroupSpec(0, ElementType.LONG, null, new BlockHash.TopNDef(0, true, false, limit))),
+                blockFactory,
+                32 * 1024,
+                forcePackedHash,
+                randomIntBetween(1, limit)
+            )
+        ) {
+            assertThat(hash, not(instanceOf(LongTopNBlockHash.class)));
+        }
+    }
+
     /**
      * Hash some values into a single block of group ids. If the hash produces
      * more than one block of group ids this will fail.
@@ -1713,6 +1744,6 @@ private BlockHash buildBlockHash(int emitBatchSize, Block... values) {
         }
         return forcePackedHash
             ? new PackedValuesBlockHash(specs, blockFactory, emitBatchSize)
-            : BlockHash.build(specs, blockFactory, emitBatchSize, true);
+            : BlockHash.build(specs, blockFactory, emitBatchSize, true, 100);
     }
 }

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/blockhash/CategorizeBlockHashTests.java

@@ -513,7 +513,8 @@ public void testCategorize_withDriver() {
                         new MaxLongAggregatorFunctionSupplier().groupingAggregatorFactory(AggregatorMode.INITIAL, List.of(1))
                     ),
                     16 * 1024,
-                    analysisRegistry
+                    analysisRegistry,
+                    100
                 ).get(driverContext)
             ),
             new PageConsumerOperator(intermediateOutput::add)
@@ -532,7 +533,8 @@ public void testCategorize_withDriver() {
                         new MaxLongAggregatorFunctionSupplier().groupingAggregatorFactory(AggregatorMode.INITIAL, List.of(1))
                     ),
                     16 * 1024,
-                    analysisRegistry
+                    analysisRegistry,
+                    100
                 ).get(driverContext)
             ),
             new PageConsumerOperator(intermediateOutput::add)
@@ -553,7 +555,8 @@ public void testCategorize_withDriver() {
                         new MaxLongAggregatorFunctionSupplier().groupingAggregatorFactory(AggregatorMode.FINAL, List.of(3, 4))
                     ),
                     16 * 1024,
-                    analysisRegistry
+                    analysisRegistry,
+                    100
                 ).get(driverContext)
             ),
             new PageConsumerOperator(finalOutput::add)

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/blockhash/CategorizePackedValuesBlockHashTests.java

@@ -151,7 +151,8 @@ public void testCategorize_withDriver() {
                     AggregatorMode.INITIAL,
                     List.of(new ValuesBytesRefAggregatorFunctionSupplier().groupingAggregatorFactory(AggregatorMode.INITIAL, List.of(0))),
                     16 * 1024,
-                    analysisRegistry
+                    analysisRegistry,
+                    100
                 ).get(driverContext)
             ),
             new PageConsumerOperator(intermediateOutput::add)
@@ -167,7 +168,8 @@ public void testCategorize_withDriver() {
                     AggregatorMode.INITIAL,
                     List.of(new ValuesBytesRefAggregatorFunctionSupplier().groupingAggregatorFactory(AggregatorMode.INITIAL, List.of(0))),
                     16 * 1024,
-                    analysisRegistry
+                    analysisRegistry,
+                    100
                 ).get(driverContext)
             ),
             new PageConsumerOperator(intermediateOutput::add)
@@ -185,7 +187,8 @@ public void testCategorize_withDriver() {
                     AggregatorMode.FINAL,
                     List.of(new ValuesBytesRefAggregatorFunctionSupplier().groupingAggregatorFactory(AggregatorMode.FINAL, List.of(2))),
                     16 * 1024,
-                    analysisRegistry
+                    analysisRegistry,
+                    100
                 ).get(driverContext)
             ),
             new PageConsumerOperator(finalOutput::add)

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/HashAggregationOperatorTests.java

@@ -64,7 +64,8 @@ protected Operator.OperatorFactory simpleWithMode(SimpleOptions options, Aggrega
                 new MaxLongAggregatorFunctionSupplier().groupingAggregatorFactory(mode, maxChannels)
             ),
             randomPageSize(),
-            null
+            null,
+            100
         );
     }
 
@@ -103,6 +104,7 @@ protected void assertSimpleOutput(List<Page> input, List<Page> results) {
 
     public void testTopNNullsLast() {
         boolean ascOrder = randomBoolean();
+        int limit = 3;
         var groups = new Long[] { 0L, 10L, 20L, 30L, 40L, 50L };
         if (ascOrder) {
             Arrays.sort(groups, Comparator.reverseOrder());
@@ -113,14 +115,15 @@ public void testTopNNullsLast() {
 
         try (
             var operator = new HashAggregationOperator.HashAggregationOperatorFactory(
-                List.of(new BlockHash.GroupSpec(groupChannel, ElementType.LONG, null, new BlockHash.TopNDef(0, ascOrder, false, 3))),
+                List.of(new BlockHash.GroupSpec(groupChannel, ElementType.LONG, null, new BlockHash.TopNDef(0, ascOrder, false, limit))),
                 mode,
                 List.of(
                     new SumLongAggregatorFunctionSupplier().groupingAggregatorFactory(mode, aggregatorChannels),
                     new MaxLongAggregatorFunctionSupplier().groupingAggregatorFactory(mode, aggregatorChannels)
                 ),
                 randomPageSize(),
-                null
+                null,
+                randomIntBetween(limit, 1000)
             ).get(driverContext())
         ) {
             var page = new Page(
@@ -180,6 +183,7 @@ public void testTopNNullsLast() {
 
     public void testTopNNullsFirst() {
         boolean ascOrder = randomBoolean();
+        int limit = 3;
         var groups = new Long[] { 0L, 10L, 20L, 30L, 40L, 50L };
         if (ascOrder) {
             Arrays.sort(groups, Comparator.reverseOrder());
@@ -190,14 +194,15 @@ public void testTopNNullsFirst() {
 
         try (
             var operator = new HashAggregationOperator.HashAggregationOperatorFactory(
-                List.of(new BlockHash.GroupSpec(groupChannel, ElementType.LONG, null, new BlockHash.TopNDef(0, ascOrder, true, 3))),
+                List.of(new BlockHash.GroupSpec(groupChannel, ElementType.LONG, null, new BlockHash.TopNDef(0, ascOrder, true, limit))),
                 mode,
                 List.of(
                     new SumLongAggregatorFunctionSupplier().groupingAggregatorFactory(mode, aggregatorChannels),
                     new MaxLongAggregatorFunctionSupplier().groupingAggregatorFactory(mode, aggregatorChannels)
                 ),
                 randomPageSize(),
-                null
+                null,
+                randomIntBetween(limit, 1000)
             ).get(driverContext())
         ) {
             var page = new Page(

x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/EsqlQueryGenerator.java

@@ -23,6 +23,7 @@
 import org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe.RenameGenerator;
 import org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe.SortGenerator;
 import org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe.StatsGenerator;
+import org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe.TopNStatsGenerator;
 import org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe.WhereGenerator;
 import org.elasticsearch.xpack.esql.qa.rest.generative.command.source.FromGenerator;
 
@@ -65,6 +66,7 @@ public record QueryExecuted(String query, int depth, List<Column> outputSchema,
         RenameGenerator.INSTANCE,
         SortGenerator.INSTANCE,
         StatsGenerator.INSTANCE,
+        TopNStatsGenerator.INSTANCE,
         WhereGenerator.INSTANCE
     );