Chunk hash aggregation output

Author: dnhatn

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

@@ -181,9 +181,11 @@ private static Operator operator(DriverContext driverContext, String grouping, S
             );
             default -> throw new IllegalArgumentException("unsupported grouping [" + grouping + "]");
         };
+        int pageSize = 16 * 1024;
         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(), pageSize, false),
+            pageSize,
             driverContext
         );
     }

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

@@ -22,6 +22,7 @@
 import org.elasticsearch.compute.data.IntBlock;
 import org.elasticsearch.compute.data.IntVector;
 import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.index.analysis.AnalysisRegistry;
@@ -58,12 +59,14 @@ public Operator get(DriverContext driverContext) {
                         analysisRegistry,
                         maxPageSize
                     ),
+                    maxPageSize,
                     driverContext
                 );
             }
             return new HashAggregationOperator(
                 aggregators,
                 () -> BlockHash.build(groups, driverContext.blockFactory(), maxPageSize, false),
+                maxPageSize,
                 driverContext
             );
         }
@@ -78,9 +81,10 @@ public String describe() {
         }
     }
 
-    private boolean finished;
-    private Page output;
+    private final int maxPageSize;
+    private Emitter emitter;
 
+    private boolean blockHashClosed = false;
     private final BlockHash blockHash;
 
     private final List<GroupingAggregator> aggregators;
@@ -96,7 +100,7 @@ public String describe() {
      */
     private long aggregationNanos;
     /**
-     * Count of pages this operator has processed.
+     * Count of input pages this operator has processed.
      */
     private int pagesProcessed;
     /**
@@ -112,8 +116,10 @@ public String describe() {
     public HashAggregationOperator(
         List<GroupingAggregator.Factory> aggregators,
         Supplier<BlockHash> blockHash,
+        int maxPageSize,
         DriverContext driverContext
     ) {
+        this.maxPageSize = maxPageSize;
         this.aggregators = new ArrayList<>(aggregators.size());
         this.driverContext = driverContext;
         boolean success = false;
@@ -132,7 +138,7 @@ public HashAggregationOperator(
 
     @Override
     public boolean needsInput() {
-        return finished == false;
+        return emitter == null;
     }
 
     @Override
@@ -201,59 +207,97 @@ public void close() {
 
     @Override
     public Page getOutput() {
-        Page p = output;
-        if (p != null) {
-            rowsEmitted += p.getPositionCount();
+        if (emitter == null) {
+            return null;
         }
-        output = null;
-        return p;
+        return emitter.nextPage();
     }
 
-    @Override
-    public void finish() {
-        if (finished) {
-            return;
+    private class Emitter implements Releasable {
+        private final int[] aggBlockCounts;
+        private int position = -1;
+        private IntVector allSelected = null;
+        private Block[] allKeys;
+
+        Emitter(int[] aggBlockCounts) {
+            this.aggBlockCounts = aggBlockCounts;
         }
-        finished = true;
-        Block[] blocks = null;
-        IntVector selected = null;
-        boolean success = false;
-        try {
-            selected = blockHash.nonEmpty();
-            Block[] keys = blockHash.getKeys();
-            int[] aggBlockCounts = aggregators.stream().mapToInt(GroupingAggregator::evaluateBlockCount).toArray();
-            blocks = new Block[keys.length + Arrays.stream(aggBlockCounts).sum()];
-            System.arraycopy(keys, 0, blocks, 0, keys.length);
-            int offset = keys.length;
-            for (int i = 0; i < aggregators.size(); i++) {
-                var aggregator = aggregators.get(i);
-                aggregator.evaluate(blocks, offset, selected, driverContext);
-                offset += aggBlockCounts[i];
+
+        Page nextPage() {
+            if (position == -1) {
+                position = 0;
+                // TODO: chunk selected and keys
+                allKeys = blockHash.getKeys();
+                allSelected = blockHash.nonEmpty();
+                blockHashClosed = true;
+                blockHash.close();
             }
-            output = new Page(blocks);
-            success = true;
-        } finally {
-            // selected should always be closed
-            if (selected != null) {
-                selected.close();
+            final int endPosition = Math.toIntExact(Math.min(position + (long) maxPageSize, allSelected.getPositionCount()));
+            if (endPosition == position) {
+                return null;
             }
-            if (success == false && blocks != null) {
-                Releasables.closeExpectNoException(blocks);
+            final boolean singlePage = position == 0 && endPosition == allSelected.getPositionCount();
+            final Block[] blocks = new Block[allKeys.length + Arrays.stream(aggBlockCounts).sum()];
+            IntVector selected = null;
+            boolean success = false;
+            try {
+                if (singlePage) {
+                    this.allSelected.incRef();
+                    selected = this.allSelected;
+                    for (int i = 0; i < allKeys.length; i++) {
+                        allKeys[i].incRef();
+                        blocks[i] = allKeys[i];
+                    }
+                } else {
+                    final int[] positions = new int[endPosition - position];
+                    for (int i = 0; i < positions.length; i++) {
+                        positions[i] = position + i;
+                    }
+                    selected = allSelected.filter(positions);
+                    for (int keyIndex = 0; keyIndex < allKeys.length; keyIndex++) {
+                        blocks[keyIndex] = allKeys[keyIndex].filter(positions);
+                    }
+                }
+                int blockOffset = allKeys.length;
+                for (int i = 0; i < aggregators.size(); i++) {
+                    aggregators.get(i).evaluate(blocks, blockOffset, selected, driverContext);
+                    blockOffset += aggBlockCounts[i];
+                }
+                var output = new Page(blocks);
+                rowsEmitted += output.getPositionCount();
+                success = true;
+                return output;
+            } finally {
+                position = endPosition;
+                Releasables.close(selected, success ? null : Releasables.wrap(blocks));
             }
         }
+
+        @Override
+        public void close() {
+            Releasables.close(allSelected, allKeys != null ? Releasables.wrap(allKeys) : null);
+        }
+
+        boolean doneEmitting() {
+            return allSelected != null && position >= allSelected.getPositionCount();
+        }
+    }
+
+    @Override
+    public void finish() {
+        if (emitter == null) {
+            emitter = new Emitter(aggregators.stream().mapToInt(GroupingAggregator::evaluateBlockCount).toArray());
+        }
     }
 
     @Override
     public boolean isFinished() {
-        return finished && output == null;
+        return emitter != null && emitter.doneEmitting();
     }
 
     @Override
     public void close() {
-        if (output != null) {
-            output.releaseBlocks();
-        }
-        Releasables.close(blockHash, () -> Releasables.close(aggregators));
+        Releasables.close(emitter, blockHashClosed ? null : blockHash, () -> Releasables.close(aggregators));
     }
 
     @Override
@@ -297,7 +341,7 @@ public static class Status implements Operator.Status {
          */
         private final long aggregationNanos;
         /**
-         * Count of pages this operator has processed.
+         * Count of input pages this operator has processed.
          */
         private final int pagesProcessed;
         /**
@@ -371,7 +415,7 @@ public long aggregationNanos() {
         }
 
         /**
-         * Count of pages this operator has processed.
+         * Count of input pages this operator has processed.
          */
         public int pagesProcessed() {
             return pagesProcessed;

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

@@ -205,16 +205,16 @@ public Page getOutput() {
             return null;
         }
         if (valuesAggregator != null) {
-            try {
-                return valuesAggregator.getOutput();
-            } finally {
-                final ValuesAggregator aggregator = this.valuesAggregator;
-                this.valuesAggregator = null;
-                Releasables.close(aggregator);
+            final Page output = valuesAggregator.getOutput();
+            if (output == null) {
+                Releasables.close(valuesAggregator, () -> this.valuesAggregator = null);
+            } else {
+                return output;
             }
         }
         if (ordinalAggregators.isEmpty() == false) {
             try {
+                // TODO: chunk output pages
                 return mergeOrdinalsSegmentResults();
             } catch (IOException e) {
                 throw new UncheckedIOException(e);
@@ -510,6 +510,7 @@ private static class ValuesAggregator implements Releasable {
                     maxPageSize,
                     false
                 ),
+                maxPageSize,
                 driverContext
             );
         }

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

@@ -64,6 +64,7 @@ public Operator get(DriverContext driverContext) {
             return new HashAggregationOperator(
                 aggregators,
                 () -> new TimeSeriesBlockHash(tsHashChannel, timeBucketChannel, driverContext),
+                maxPageSize,
                 driverContext
             );
         }
@@ -99,6 +100,7 @@ public Operator get(DriverContext driverContext) {
             return new HashAggregationOperator(
                 aggregators,
                 () -> BlockHash.build(hashGroups, driverContext.blockFactory(), maxPageSize, false),
+                maxPageSize,
                 driverContext
             );
         }
@@ -127,6 +129,7 @@ public Operator get(DriverContext driverContext) {
             return new HashAggregationOperator(
                 aggregators,
                 () -> BlockHash.build(groupings, driverContext.blockFactory(), maxPageSize, false),
+                maxPageSize,
                 driverContext
             );
         }

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

@@ -211,6 +211,7 @@ public String toString() {
                             randomPageSize(),
                             false
                         ),
+                        randomPageSize(),
                         driverContext
                     )
                 );