Release grouped top n memory on output

core/trino-main/src/main/java/io/trino/operator/GroupedTopNRankBuilder.java

@@ -20,6 +20,7 @@
 import io.trino.spi.Page;
 import io.trino.spi.PageBuilder;
 import io.trino.spi.type.Type;
+import jakarta.annotation.Nullable;
 
 import java.util.Iterator;
 import java.util.List;
@@ -42,7 +43,8 @@ public class GroupedTopNRankBuilder
     private final List<Type> sourceTypes;
     private final boolean produceRanking;
     private final int[] groupByChannels;
-    private final GroupByHash groupByHash;
+    @Nullable
+    private GroupByHash groupByHash; // null after output starts
     private final PageWithPositionComparator comparator;
     private final RowReferencePageManager pageManager = new RowReferencePageManager();
     private final GroupedTopNRankAccumulator groupedTopNRankAccumulator;
@@ -102,6 +104,9 @@ public long hashCode(long rowId)
     @Override
     public Work<?> processPage(Page page)
     {
+        if (groupByHash == null) {
+            throw new IllegalStateException("already producing results");
+        }
         return new TransformWork<>(
                 groupByHash.getGroupIds(page.getColumns(groupByChannels)),
                 groupIds -> {
@@ -113,14 +118,19 @@ public Work<?> processPage(Page page)
     @Override
     public Iterator<Page> buildResult()
     {
-        return new ResultIterator();
+        if (groupByHash == null) {
+            throw new IllegalStateException("already producing results");
+        }
+        int groupIdCount = groupByHash.getGroupCount();
+        groupByHash = null;
+        return new ResultIterator(groupIdCount);
     }
 
     @Override
     public long getEstimatedSizeInBytes()
     {
         return INSTANCE_SIZE
-                + groupByHash.getEstimatedSize()
+                + (groupByHash == null ? 0L : groupByHash.getEstimatedSize())
                 + pageManager.sizeOf()
                 + groupedTopNRankAccumulator.sizeOf();
     }
@@ -148,15 +158,16 @@ private class ResultIterator
             extends AbstractIterator<Page>
     {
         private final PageBuilder pageBuilder;
-        private final int groupIdCount = groupByHash.getGroupCount();
+        private final int groupIdCount;
         private int currentGroupId = -1;
         private final LongBigArray rowIdOutput = new LongBigArray();
         private final LongBigArray rankingOutput = new LongBigArray();
         private long currentGroupSize;
         private int currentIndexInGroup;
 
-        ResultIterator()
+        private ResultIterator(int groupIdCount)
         {
+            this.groupIdCount = groupIdCount;
             ImmutableList.Builder<Type> sourceTypesBuilders = ImmutableList.<Type>builder().addAll(sourceTypes);
             if (produceRanking) {
                 sourceTypesBuilders.add(BIGINT);

core/trino-main/src/main/java/io/trino/operator/GroupedTopNRowNumberBuilder.java

@@ -21,6 +21,7 @@
 import io.trino.spi.Page;
 import io.trino.spi.PageBuilder;
 import io.trino.spi.type.Type;
+import jakarta.annotation.Nullable;
 
 import java.util.Iterator;
 import java.util.List;
@@ -42,7 +43,8 @@ public class GroupedTopNRowNumberBuilder
     private final List<Type> sourceTypes;
     private final boolean produceRowNumber;
     private final int[] groupByChannels;
-    private final GroupByHash groupByHash;
+    @Nullable
+    private GroupByHash groupByHash; // null after output starts
     private final RowReferencePageManager pageManager = new RowReferencePageManager();
     private final GroupedTopNRowNumberAccumulator groupedTopNRowNumberAccumulator;
     private final PageWithPositionComparator comparator;
@@ -77,6 +79,9 @@ public GroupedTopNRowNumberBuilder(
     @Override
     public Work<?> processPage(Page page)
     {
+        if (groupByHash == null) {
+            throw new IllegalStateException("already producing output");
+        }
         return new TransformWork<>(
                 groupByHash.getGroupIds(page.getColumns(groupByChannels)),
                 groupIds -> {
@@ -88,14 +93,19 @@ public Work<?> processPage(Page page)
     @Override
     public Iterator<Page> buildResult()
     {
-        return new ResultIterator();
+        if (groupByHash == null) {
+            throw new IllegalStateException("already producing output");
+        }
+        int groupIdCount = groupByHash.getGroupCount();
+        groupByHash = null;
+        return new ResultIterator(groupIdCount);
     }
 
     @Override
     public long getEstimatedSizeInBytes()
     {
         return INSTANCE_SIZE
-                + groupByHash.getEstimatedSize()
+                + (groupByHash == null ? 0L : groupByHash.getEstimatedSize())
                 + pageManager.sizeOf()
                 + groupedTopNRowNumberAccumulator.sizeOf();
     }
@@ -119,6 +129,7 @@ private void processPage(Page newPage, int groupCount, int[] groupIds)
         pageManager.compactIfNeeded();
     }
 
+    @Nullable
     @VisibleForTesting
     GroupByHash getGroupByHash()
     {
@@ -129,14 +140,15 @@ private class ResultIterator
             extends AbstractIterator<Page>
     {
         private final PageBuilder pageBuilder;
-        private final int groupIdCount = groupByHash.getGroupCount();
+        private final int groupIdCount;
         private int currentGroupId = -1;
         private final LongBigArray rowIdOutput = new LongBigArray();
         private long currentGroupSize;
         private int currentIndexInGroup;
 
-        ResultIterator()
+        private ResultIterator(int groupIdCount)
         {
+            this.groupIdCount = groupIdCount;
             ImmutableList.Builder<Type> sourceTypesBuilders = ImmutableList.<Type>builder().addAll(sourceTypes);
             if (produceRowNumber) {
                 sourceTypesBuilders.add(BIGINT);