Cleanup DistinctWindowAccumulator

Use exact distinct positions count instead of estimating the size, clean
up control flow to handle no distinct positions earlier.

Author: pettyjamesm

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

@@ -42,6 +42,11 @@ private MarkDistinctHash(MarkDistinctHash other)
         nextDistinctId = other.nextDistinctId;
     }
 
+    public long getGroupCount()
+    {
+        return groupByHash.getGroupCount();
+    }
+
     public long getEstimatedSize()
     {
         return groupByHash.getEstimatedSize();

core/trino-main/src/main/java/io/trino/operator/aggregation/DistinctWindowAccumulator.java

@@ -30,14 +30,13 @@
 import io.trino.spi.function.WindowAccumulator;
 import io.trino.spi.function.WindowIndex;
 import io.trino.spi.type.Type;
-import it.unimi.dsi.fastutil.ints.IntArrayList;
 
 import java.util.List;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkState;
 import static io.trino.spi.type.BooleanType.BOOLEAN;
-import static java.lang.Math.min;
+import static java.lang.Math.toIntExact;
 import static java.util.Objects.requireNonNull;
 
 public class DistinctWindowAccumulator
@@ -128,39 +127,46 @@ public void addInput(WindowIndex index, int startPosition, int endPosition)
 
     private void indexCurrentPage(Page page)
     {
+        long initialGroupCount = hash.getGroupCount();
         Work<Block> work = hash.markDistinctRows(page);
         checkState(work.process());
         Block distinctMask = work.getResult();
 
         int positionCount = distinctMask.getPositionCount();
         checkArgument(positionCount == page.getPositionCount(), "Page position count does not match distinct mask position count");
-        PagesIndex pagesIndex = pagesIndexFactory.newPagesIndex(argumentTypes, positionCount);
+
+        int distinctPositions = toIntExact(hash.getGroupCount() - initialGroupCount);
+        if (distinctPositions == 0) {
+            return;
+        }
+        PagesIndex pagesIndex = pagesIndexFactory.newPagesIndex(argumentTypes, distinctPositions);
 
         if (distinctMask instanceof RunLengthEncodedBlock) {
-            if (test(distinctMask, 0)) {
-                // all positions selected
-                pagesIndex.addPage(page);
-            }
+            // all positions selected
+            checkState(test(distinctMask, 0), "all positions must be distinct");
+            pagesIndex.addPage(page);
         }
         else {
-            IntArrayList selectedPositions = new IntArrayList(min(128, positionCount));
+            int[] selectedPositions = new int[distinctPositions];
+            int selectedIndex = 0;
             for (int position = 0; position < positionCount; position++) {
                 if (test(distinctMask, position)) {
-                    selectedPositions.add(position);
+                    selectedPositions[selectedIndex++] = position;
                 }
             }
+            checkState(selectedIndex == selectedPositions.length, "Invalid positions in distinct mask");
 
             Block[] filteredBlocks = new Block[argumentChannels.size()];
             for (int channel = 0; channel < argumentChannels.size(); channel++) {
-                filteredBlocks[channel] = page.getBlock(channel).copyPositions(selectedPositions.elements(), 0, selectedPositions.size());
+                filteredBlocks[channel] = page.getBlock(channel).copyPositions(selectedPositions, 0, selectedPositions.length);
             }
-            pagesIndex.addPage(new Page(selectedPositions.size(), filteredBlocks));
+            pagesIndex.addPage(new Page(selectedPositions.length, filteredBlocks));
         }
         int selectedPositionsCount = pagesIndex.getPositionCount();
-        if (selectedPositionsCount > 0) {
-            PagesWindowIndex selectedWindowIndex = new PagesWindowIndex(pagesIndex, 0, selectedPositionsCount);
-            delegate.addInput(selectedWindowIndex, 0, selectedPositionsCount - 1);
-        }
+        checkState(selectedPositionsCount == distinctPositions, "unexpected pagesIndex positions: %s <> %s", selectedPositionsCount, distinctPositions);
+
+        PagesWindowIndex selectedWindowIndex = new PagesWindowIndex(pagesIndex, 0, selectedPositionsCount);
+        delegate.addInput(selectedWindowIndex, 0, selectedPositionsCount - 1);
     }
 
     private static boolean test(Block block, int position)