ESQL: Skip unused STATS groups by adding a Top N BlockHash implementation

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

@@ -73,6 +73,7 @@ public class AggregatorBenchmark {
     static final int BLOCK_LENGTH = 8 * 1024;
     private static final int OP_COUNT = 1024;
     private static final int GROUPS = 5;
+    private static final int TOP_N_LIMIT = 3;
 
     private static final BlockFactory blockFactory = BlockFactory.getInstance(
         new NoopCircuitBreaker("noop"),
@@ -90,6 +91,7 @@ public class AggregatorBenchmark {
     private static final String TWO_ORDINALS = "two_" + ORDINALS;
     private static final String LONGS_AND_BYTES_REFS = LONGS + "_and_" + BYTES_REFS;
     private static final String TWO_LONGS_AND_BYTES_REFS = "two_" + LONGS + "_and_" + BYTES_REFS;
+    private static final String TOP_N_LONGS = "top_n_" + LONGS;
 
     private static final String VECTOR_DOUBLES = "vector_doubles";
     private static final String HALF_NULL_DOUBLES = "half_null_doubles";
@@ -147,7 +149,8 @@ static void selfTest() {
             TWO_BYTES_REFS,
             TWO_ORDINALS,
             LONGS_AND_BYTES_REFS,
-            TWO_LONGS_AND_BYTES_REFS }
+            TWO_LONGS_AND_BYTES_REFS,
+            TOP_N_LONGS }
     )
     public String grouping;
 
@@ -161,8 +164,7 @@ static void selfTest() {
     public String filter;
 
     private static Operator operator(DriverContext driverContext, String grouping, String op, String dataType, String filter) {
-
-        if (grouping.equals("none")) {
+        if (grouping.equals(NONE)) {
             return new AggregationOperator(
                 List.of(supplier(op, dataType, filter).aggregatorFactory(AggregatorMode.SINGLE, List.of(0)).apply(driverContext)),
                 driverContext
@@ -188,6 +190,9 @@ private static Operator operator(DriverContext driverContext, String grouping, S
                 new BlockHash.GroupSpec(1, ElementType.LONG),
                 new BlockHash.GroupSpec(2, ElementType.BYTES_REF)
             );
+            case TOP_N_LONGS -> List.of(
+                new BlockHash.GroupSpec(0, ElementType.LONG, false, new BlockHash.TopNDef(0, true, true, TOP_N_LIMIT))
+            );
             default -> throw new IllegalArgumentException("unsupported grouping [" + grouping + "]");
         };
         return new HashAggregationOperator(
@@ -271,10 +276,14 @@ private static void checkGrouped(String prefix, String grouping, String op, Stri
             case BOOLEANS -> 2;
             default -> GROUPS;
         };
+        int availableGroups = switch (grouping) {
+            case TOP_N_LONGS -> TOP_N_LIMIT;
+            default -> groups;
+        };
         switch (op) {
             case AVG -> {
                 DoubleBlock dValues = (DoubleBlock) values;
-                for (int g = 0; g < groups; g++) {
+                for (int g = 0; g < availableGroups; g++) {
                     long group = g;
                     long sum = LongStream.range(0, BLOCK_LENGTH).filter(l -> l % groups == group).sum();
                     long count = LongStream.range(0, BLOCK_LENGTH).filter(l -> l % groups == group).count();
@@ -286,7 +295,7 @@ private static void checkGrouped(String prefix, String grouping, String op, Stri
             }
             case COUNT -> {
                 LongBlock lValues = (LongBlock) values;
-                for (int g = 0; g < groups; g++) {
+                for (int g = 0; g < availableGroups; g++) {
                     long group = g;
                     long expected = LongStream.range(0, BLOCK_LENGTH).filter(l -> l % groups == group).count() * opCount;
                     if (lValues.getLong(g) != expected) {
@@ -296,7 +305,7 @@ private static void checkGrouped(String prefix, String grouping, String op, Stri
             }
             case COUNT_DISTINCT -> {
                 LongBlock lValues = (LongBlock) values;
-                for (int g = 0; g < groups; g++) {
+                for (int g = 0; g < availableGroups; g++) {
                     long group = g;
                     long expected = LongStream.range(0, BLOCK_LENGTH).filter(l -> l % groups == group).distinct().count();
                     long count = lValues.getLong(g);
@@ -310,15 +319,15 @@ private static void checkGrouped(String prefix, String grouping, String op, Stri
                 switch (dataType) {
                     case LONGS -> {
                         LongBlock lValues = (LongBlock) values;
-                        for (int g = 0; g < groups; g++) {
+                        for (int g = 0; g < availableGroups; g++) {
                             if (lValues.getLong(g) != (long) g) {
                                 throw new AssertionError(prefix + "expected [" + g + "] but was [" + lValues.getLong(g) + "]");
                             }
                         }
                     }
                     case DOUBLES -> {
                         DoubleBlock dValues = (DoubleBlock) values;
-                        for (int g = 0; g < groups; g++) {
+                        for (int g = 0; g < availableGroups; g++) {
                             if (dValues.getDouble(g) != (long) g) {
                                 throw new AssertionError(prefix + "expected [" + g + "] but was [" + dValues.getDouble(g) + "]");
                             }
@@ -331,7 +340,7 @@ private static void checkGrouped(String prefix, String grouping, String op, Stri
                 switch (dataType) {
                     case LONGS -> {
                         LongBlock lValues = (LongBlock) values;
-                        for (int g = 0; g < groups; g++) {
+                        for (int g = 0; g < availableGroups; g++) {
                             long group = g;
                             long expected = LongStream.range(0, BLOCK_LENGTH).filter(l -> l % groups == group).max().getAsLong();
                             if (lValues.getLong(g) != expected) {
@@ -341,7 +350,7 @@ private static void checkGrouped(String prefix, String grouping, String op, Stri
                     }
                     case DOUBLES -> {
                         DoubleBlock dValues = (DoubleBlock) values;
-                        for (int g = 0; g < groups; g++) {
+                        for (int g = 0; g < availableGroups; g++) {
                             long group = g;
                             long expected = LongStream.range(0, BLOCK_LENGTH).filter(l -> l % groups == group).max().getAsLong();
                             if (dValues.getDouble(g) != expected) {
@@ -356,7 +365,7 @@ private static void checkGrouped(String prefix, String grouping, String op, Stri
                 switch (dataType) {
                     case LONGS -> {
                         LongBlock lValues = (LongBlock) values;
-                        for (int g = 0; g < groups; g++) {
+                        for (int g = 0; g < availableGroups; g++) {
                             long group = g;
                             long expected = LongStream.range(0, BLOCK_LENGTH).filter(l -> l % groups == group).sum() * opCount;
                             if (lValues.getLong(g) != expected) {
@@ -366,7 +375,7 @@ private static void checkGrouped(String prefix, String grouping, String op, Stri
                     }
                     case DOUBLES -> {
                         DoubleBlock dValues = (DoubleBlock) values;
-                        for (int g = 0; g < groups; g++) {
+                        for (int g = 0; g < availableGroups; g++) {
                             long group = g;
                             long expected = LongStream.range(0, BLOCK_LENGTH).filter(l -> l % groups == group).sum() * opCount;
                             if (dValues.getDouble(g) != expected) {
@@ -391,6 +400,14 @@ private static void checkGroupingBlock(String prefix, String grouping, Block blo
                     }
                 }
             }
+            case TOP_N_LONGS -> {
+                LongBlock groups = (LongBlock) block;
+                for (int g = 0; g < TOP_N_LIMIT; g++) {
+                    if (groups.getLong(g) != (long) g) {
+                        throw new AssertionError(prefix + "bad group expected [" + g + "] but was [" + groups.getLong(g) + "]");
+                    }
+                }
+            }
             case INTS -> {
                 IntBlock groups = (IntBlock) block;
                 for (int g = 0; g < GROUPS; g++) {
@@ -495,7 +512,7 @@ private static void checkUngrouped(String prefix, String op, String dataType, Pa
 
     private static Page page(BlockFactory blockFactory, String grouping, String blockType) {
         Block dataBlock = dataBlock(blockFactory, blockType);
-        if (grouping.equals("none")) {
+        if (grouping.equals(NONE)) {
             return new Page(dataBlock);
         }
         List<Block> blocks = groupingBlocks(grouping, blockType);
@@ -564,7 +581,7 @@ private static Block groupingBlock(String grouping, String blockType) {
             default -> throw new UnsupportedOperationException("bad grouping [" + grouping + "]");
         };
         return switch (grouping) {
-            case LONGS -> {
+            case TOP_N_LONGS, LONGS -> {
                 var builder = blockFactory.newLongBlockBuilder(BLOCK_LENGTH);
                 for (int i = 0; i < BLOCK_LENGTH; i++) {
                     for (int v = 0; v < valuesPerGroup; v++) {

docs/changelog/127148.yaml

@@ -0,0 +1,5 @@
+pr: 127148
+summary: Skip unused STATS groups by adding a Top N `BlockHash` implementation
+area: ES|QL
+type: enhancement
+issues: []

server/src/main/java/org/elasticsearch/common/util/BinarySearcher.java

@@ -38,7 +38,7 @@ public abstract class BinarySearcher {
     /**
      * @return the index who's underlying value is closest to the value being searched for.
      */
-    private int getClosestIndex(int index1, int index2) {
+    protected int getClosestIndex(int index1, int index2) {
         if (distance(index1) < distance(index2)) {
             return index1;
         } else {

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

@@ -23,6 +23,7 @@
 import org.elasticsearch.compute.data.IntBlock;
 import org.elasticsearch.compute.data.IntVector;
 import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.ReleasableIterator;
 import org.elasticsearch.index.analysis.AnalysisRegistry;
@@ -113,13 +114,30 @@ public abstract class BlockHash implements Releasable, SeenGroupIds {
     @Override
     public abstract BitArray seenGroupIds(BigArrays bigArrays);
 
+    /**
+     * Configuration for a BlockHash group spec that is later sorted and limited (Top-N).
+     * <p>
+     *     Part of a performance improvement to avoid aggregating groups that will not be used.
+     * </p>
+     *
+     * @param order The order of this group in the sort, starting at 0
+     * @param asc True if this group will be sorted ascending. False if descending.
+     * @param nullsFirst True if the nulls should be the first elements in the TopN. False if they should be kept last.
+     * @param limit The number of elements to keep, including nulls.
+     */
+    public record TopNDef(int order, boolean asc, boolean nullsFirst, int limit) {}
+
     /**
      * @param isCategorize Whether this group is a CATEGORIZE() or not.
      *                     May be changed in the future when more stateful grouping functions are added.
      */
-    public record GroupSpec(int channel, ElementType elementType, boolean isCategorize) {
+    public record GroupSpec(int channel, ElementType elementType, boolean isCategorize, @Nullable TopNDef topNDef) {
         public GroupSpec(int channel, ElementType elementType) {
-            this(channel, elementType, false);
+            this(channel, elementType, false, null);
+        }
+
+        public GroupSpec(int channel, ElementType elementType, boolean isCategorize) {
+            this(channel, elementType, isCategorize, null);
         }
     }
 
@@ -134,7 +152,12 @@ public GroupSpec(int channel, ElementType elementType) {
      */
     public static BlockHash build(List<GroupSpec> groups, BlockFactory blockFactory, int emitBatchSize, boolean allowBrokenOptimizations) {
         if (groups.size() == 1) {
-            return newForElementType(groups.get(0).channel(), groups.get(0).elementType(), blockFactory);
+            GroupSpec group = groups.get(0);
+            if (group.topNDef() != null && group.elementType() == ElementType.LONG) {
+                TopNDef topNDef = group.topNDef();
+                return new LongTopNBlockHash(group.channel(), topNDef.asc(), topNDef.nullsFirst(), topNDef.limit(), blockFactory);
+            }
+            return newForElementType(group.channel(), group.elementType(), blockFactory);
         }
         if (groups.stream().allMatch(g -> g.elementType == ElementType.BYTES_REF)) {
             switch (groups.size()) {