Move compactRelatedBlocks from Page to DictionaryBlock

trinodb · Sep 18, 2022 · 761ffbd · 761ffbd
1 parent 08342e5
commit 761ffbd
Show file tree

Hide file tree

Showing 2 changed files with 68 additions and 65 deletions.
diff --git a/core/trino-spi/src/main/java/io/trino/spi/Page.java b/core/trino-spi/src/main/java/io/trino/spi/Page.java
@@ -25,8 +25,6 @@
 import java.util.Map;
 
 import static io.airlift.slice.SizeOf.sizeOf;
-import static io.trino.spi.block.DictionaryId.randomDictionaryId;
-import static java.lang.Math.min;
 import static java.lang.String.format;
 import static java.util.Objects.requireNonNull;
 
@@ -187,7 +185,7 @@ public void compact()
 
         Map<DictionaryId, DictionaryBlockIndexes> dictionaryBlocks = getRelatedDictionaryBlocks();
         for (DictionaryBlockIndexes blockIndexes : dictionaryBlocks.values()) {
-            List<DictionaryBlock> compactBlocks = compactRelatedBlocks(blockIndexes.getBlocks());
+            List<DictionaryBlock> compactBlocks = DictionaryBlock.compactRelatedBlocks(blockIndexes.getBlocks());
             List<Integer> indexes = blockIndexes.getIndexes();
             for (int i = 0; i < compactBlocks.size(); i++) {
                 blocks[indexes.get(i)] = compactBlocks.get(i);
@@ -212,68 +210,6 @@ private Map<DictionaryId, DictionaryBlockIndexes> getRelatedDictionaryBlocks()
         return relatedDictionaryBlocks;
     }
 
-    private static List<DictionaryBlock> compactRelatedBlocks(List<DictionaryBlock> blocks)
-    {
-        DictionaryBlock firstDictionaryBlock = blocks.get(0);
-        Block dictionary = firstDictionaryBlock.getDictionary();
-
-        int positionCount = firstDictionaryBlock.getPositionCount();
-        int dictionarySize = dictionary.getPositionCount();
-
-        // determine which dictionary entries are referenced and build a reindex for them
-        int[] dictionaryPositionsToCopy = new int[min(dictionarySize, positionCount)];
-        int[] remapIndex = new int[dictionarySize];
-        Arrays.fill(remapIndex, -1);
-
-        int numberOfIndexes = 0;
-        for (int i = 0; i < positionCount; i++) {
-            int position = firstDictionaryBlock.getId(i);
-            if (remapIndex[position] == -1) {
-                dictionaryPositionsToCopy[numberOfIndexes] = position;
-                remapIndex[position] = numberOfIndexes;
-                numberOfIndexes++;
-            }
-        }
-
-        // entire dictionary is referenced
-        if (numberOfIndexes == dictionarySize) {
-            return blocks;
-        }
-
-        // compact the dictionaries
-        int[] newIds = getNewIds(positionCount, firstDictionaryBlock, remapIndex);
-        List<DictionaryBlock> outputDictionaryBlocks = new ArrayList<>(blocks.size());
-        DictionaryId newDictionaryId = randomDictionaryId();
-        for (DictionaryBlock dictionaryBlock : blocks) {
-            if (!firstDictionaryBlock.getDictionarySourceId().equals(dictionaryBlock.getDictionarySourceId())) {
-                throw new IllegalArgumentException("dictionarySourceIds must be the same");
-            }
-
-            try {
-                Block compactDictionary = dictionaryBlock.getDictionary().copyPositions(dictionaryPositionsToCopy, 0, numberOfIndexes);
-                outputDictionaryBlocks.add(new DictionaryBlock(positionCount, compactDictionary, newIds, !(compactDictionary instanceof DictionaryBlock), newDictionaryId));
-            }
-            catch (UnsupportedOperationException e) {
-                // ignore if copy positions is not supported for the dictionary
-                outputDictionaryBlocks.add(dictionaryBlock);
-            }
-        }
-        return outputDictionaryBlocks;
-    }
-
-    private static int[] getNewIds(int positionCount, DictionaryBlock dictionaryBlock, int[] remapIndex)
-    {
-        int[] newIds = new int[positionCount];
-        for (int i = 0; i < positionCount; i++) {
-            int newId = remapIndex[dictionaryBlock.getId(i)];
-            if (newId == -1) {
-                throw new IllegalStateException("reference to a non-existent key");
-            }
-            newIds[i] = newId;
-        }
-        return newIds;
-    }
-
     /**
      * Returns a page that assures all data is in memory.
      * May return the same page if all page data is already in memory.

diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/DictionaryBlock.java b/core/trino-spi/src/main/java/io/trino/spi/block/DictionaryBlock.java
@@ -17,6 +17,7 @@
 import io.airlift.slice.Slices;
 import org.openjdk.jol.info.ClassLayout;
 
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.OptionalInt;
@@ -714,4 +715,70 @@ private DictionaryBlock unnest()
 
         return new DictionaryBlock(dictionary, ids);
     }
+
+    /**
+     * Compact the dictionary down to only the used positions for a set of
+     * blocks that have been projected from the same dictionary.
+     */
+    public static List<DictionaryBlock> compactRelatedBlocks(List<DictionaryBlock> blocks)
+    {
+        DictionaryBlock firstDictionaryBlock = blocks.get(0);
+        Block dictionary = firstDictionaryBlock.getDictionary();
+
+        int positionCount = firstDictionaryBlock.getPositionCount();
+        int dictionarySize = dictionary.getPositionCount();
+
+        // determine which dictionary entries are referenced and build a reindex for them
+        int[] dictionaryPositionsToCopy = new int[min(dictionarySize, positionCount)];
+        int[] remapIndex = new int[dictionarySize];
+        Arrays.fill(remapIndex, -1);
+
+        int numberOfIndexes = 0;
+        for (int i = 0; i < positionCount; i++) {
+            int position = firstDictionaryBlock.getId(i);
+            if (remapIndex[position] == -1) {
+                dictionaryPositionsToCopy[numberOfIndexes] = position;
+                remapIndex[position] = numberOfIndexes;
+                numberOfIndexes++;
+            }
+        }
+
+        // entire dictionary is referenced
+        if (numberOfIndexes == dictionarySize) {
+            return blocks;
+        }
+
+        // compact the dictionaries
+        int[] newIds = getNewIds(positionCount, firstDictionaryBlock, remapIndex);
+        List<DictionaryBlock> outputDictionaryBlocks = new ArrayList<>(blocks.size());
+        DictionaryId newDictionaryId = randomDictionaryId();
+        for (DictionaryBlock dictionaryBlock : blocks) {
+            if (!firstDictionaryBlock.getDictionarySourceId().equals(dictionaryBlock.getDictionarySourceId())) {
+                throw new IllegalArgumentException("dictionarySourceIds must be the same");
+            }
+
+            try {
+                Block compactDictionary = dictionaryBlock.getDictionary().copyPositions(dictionaryPositionsToCopy, 0, numberOfIndexes);
+                outputDictionaryBlocks.add(new DictionaryBlock(positionCount, compactDictionary, newIds, !(compactDictionary instanceof DictionaryBlock), newDictionaryId));
+            }
+            catch (UnsupportedOperationException e) {
+                // ignore if copy positions is not supported for the dictionary
+                outputDictionaryBlocks.add(dictionaryBlock);
+            }
+        }
+        return outputDictionaryBlocks;
+    }
+
+    private static int[] getNewIds(int positionCount, DictionaryBlock dictionaryBlock, int[] remapIndex)
+    {
+        int[] newIds = new int[positionCount];
+        for (int i = 0; i < positionCount; i++) {
+            int newId = remapIndex[dictionaryBlock.getId(i)];
+            if (newId == -1) {
+                throw new IllegalStateException("reference to a non-existent key");
+            }
+            newIds[i] = newId;
+        }
+        return newIds;
+    }
 }