Included TopNLongBlockHash into BlockHash.ubuild() and GroupSpec logics, and added it to benchmark

Author: ivancea

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,12 @@ 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 +279,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 +298,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 +308,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 +322,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 +343,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 +353,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 +368,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 +378,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 +403,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 +515,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 +584,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++) {

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,18 @@ 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()) {

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

@@ -363,7 +363,7 @@ private void hashBatchesCallbackOnLast(Consumer<OrdsAndKeys> callback, Block[]..
     private BlockHash buildBlockHash(int emitBatchSize, Block... values) {
         List<BlockHash.GroupSpec> specs = new ArrayList<>(values.length);
         for (int c = 0; c < values.length; c++) {
-            specs.add(new BlockHash.GroupSpec(c, values[c].elementType()));
+            specs.add(new BlockHash.GroupSpec(c, values[c].elementType(), false, topNDef(c)));
         }
         assert forcePackedHash == false : "Packed TopN hash not implemented yet";
         /*return forcePackedHash
@@ -386,4 +386,8 @@ private String topNParametersString(int differentValues, int unusedInsertedValue
             + ", entries="
             + Math.min(differentValues, limit + unusedInsertedValues);
     }
+
+    private BlockHash.TopNDef topNDef(int order) {
+        return new BlockHash.TopNDef(order, asc, nullsFirst, limit);
+    }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/AbstractPhysicalOperationProviders.java

@@ -360,7 +360,7 @@ BlockHash.GroupSpec toHashGroupSpec() {
                 throw new EsqlIllegalArgumentException("planned to use ordinals but tried to use the hash instead");
             }
 
-            return new BlockHash.GroupSpec(channel, elementType(), Alias.unwrap(expression) instanceof Categorize);
+            return new BlockHash.GroupSpec(channel, elementType(), Alias.unwrap(expression) instanceof Categorize, null);
         }
 
         ElementType elementType() {