Fix lazy map hashtable regression

presto-array/src/main/java/com/facebook/presto/array/ReferenceCountMap.java

@@ -13,6 +13,7 @@
  */
 package com.facebook.presto.array;
 
+import com.facebook.presto.spi.block.AbstractMapBlock;
 import com.facebook.presto.spi.block.Block;
 import io.airlift.slice.SizeOf;
 import io.airlift.slice.Slice;
@@ -90,6 +91,9 @@ else if (key instanceof Slice) {
         else if (key.getClass().isArray()) {
             extraIdentity = getLength(key);
         }
+        else if (key instanceof AbstractMapBlock.HashTables) {
+            extraIdentity = (int) ((AbstractMapBlock.HashTables) key).getRetainedSizeInBytes();
+        }
         else {
             throw new IllegalArgumentException(format("Unsupported type for %s", key));
         }

presto-main/src/test/java/com/facebook/presto/block/AbstractTestBlock.java

@@ -14,6 +14,7 @@
 package com.facebook.presto.block;
 
 import com.facebook.presto.metadata.FunctionRegistry;
+import com.facebook.presto.spi.block.AbstractMapBlock.HashTables;
 import com.facebook.presto.spi.block.Block;
 import com.facebook.presto.spi.block.BlockBuilder;
 import com.facebook.presto.spi.block.BlockBuilderStatus;
@@ -154,6 +155,9 @@ else if (type == short[].class) {
                 else if (type == DictionaryId.class) {
                     retainedSize += ClassLayout.parseClass(DictionaryId.class).instanceSize();
                 }
+                else if (type == HashTables.class) {
+                    retainedSize += ((HashTables) field.get(block)).getRetainedSizeInBytes();
+                }
                 else if (type == MethodHandle.class) {
                     // MethodHandles are only used in MapBlock/MapBlockBuilder,
                     // and they are shared among blocks created by the same MapType.

presto-main/src/test/java/com/facebook/presto/block/TestMapBlock.java

@@ -18,6 +18,7 @@
 import com.facebook.presto.spi.block.Block;
 import com.facebook.presto.spi.block.BlockBuilder;
 import com.facebook.presto.spi.block.ByteArrayBlock;
+import com.facebook.presto.spi.block.MapBlock;
 import com.facebook.presto.spi.block.MapBlockBuilder;
 import com.facebook.presto.spi.block.SingleMapBlock;
 import com.facebook.presto.spi.type.MapType;
@@ -93,20 +94,32 @@ public void testLazyHashTableBuildOverBlockRegion()
         int offset = block.getPositionCount() / 2;
         Block blockRegion = block.getRegion(offset, block.getPositionCount() - offset);
 
+        assertFalse(((MapBlock) blockRegion).isHashTablesCreated());
+        assertFalse(((MapBlock) block).isHashTablesCreated());
+
         // Lazily build the hashtables for the block region and use them to do position/value check.
         Map<String, Long>[] expectedValues = Arrays.copyOfRange(values, values.length / 2, values.length);
         assertBlock(blockRegion, () -> blockBuilder.newBlockBuilderLike(null), expectedValues);
 
+        assertTrue(((MapBlock) blockRegion).isHashTablesCreated());
+        assertTrue(((MapBlock) block).isHashTablesCreated());
+
         Map<String, Long>[] valuesWithNull = alternatingNullValues(values);
         Block blockWithNull = createBlockWithValuesFromKeyValueBlock(valuesWithNull);
 
         // Create a MapBlock that is a region of another MapBlock with null values. It doesn't have hashtables built at the time of creation.
         offset = blockWithNull.getPositionCount() / 2;
         Block blockRegionWithNull = blockWithNull.getRegion(offset, blockWithNull.getPositionCount() - offset);
 
+        assertFalse(((MapBlock) blockRegionWithNull).isHashTablesCreated());
+        assertFalse(((MapBlock) blockWithNull).isHashTablesCreated());
+
         // Lazily build the hashtables for the block region and use them to do position/value check.
         Map<String, Long>[] expectedValuesWithNull = Arrays.copyOfRange(valuesWithNull, valuesWithNull.length / 2, valuesWithNull.length);
         assertBlock(blockRegionWithNull, () -> blockBuilder.newBlockBuilderLike(null), expectedValuesWithNull);
+
+        assertTrue(((MapBlock) blockRegionWithNull).isHashTablesCreated());
+        assertTrue(((MapBlock) blockWithNull).isHashTablesCreated());
     }
 
     private Map<String, Long>[] createTestMap(int... entryCounts)

presto-spi/src/main/java/com/facebook/presto/spi/block/AbstractMapBlock.java

@@ -15,6 +15,7 @@
 package com.facebook.presto.spi.block;
 
 import com.facebook.presto.spi.type.Type;
+import org.openjdk.jol.info.ClassLayout;
 
 import javax.annotation.Nullable;
 
@@ -28,6 +29,7 @@
 import static com.facebook.presto.spi.block.BlockUtil.compactArray;
 import static com.facebook.presto.spi.block.BlockUtil.compactOffsets;
 import static com.facebook.presto.spi.block.MapBlock.createMapBlockInternal;
+import static io.airlift.slice.SizeOf.sizeOf;
 import static java.util.Objects.requireNonNull;
 
 public abstract class AbstractMapBlock
@@ -55,8 +57,7 @@ public AbstractMapBlock(Type keyType, MethodHandle keyNativeHashCode, MethodHand
 
     protected abstract Block getRawValueBlock();
 
-    @Nullable
-    protected abstract int[] getHashTables();
+    protected abstract HashTables getHashTables();
 
     /**
      * offset is entry-based, not position-based. In other words,
@@ -116,7 +117,7 @@ public Block copyPositions(int[] positions, int offset, int length)
             newPosition++;
         }
 
-        int[] hashTable = getHashTables();
+        int[] hashTable = getHashTables().getRawHashTables();
         int[] newHashTable = null;
         if (hashTable != null) {
             newHashTable = new int[newOffsets[newOffsets.length - 1] * HASH_MULTIPLIER];
@@ -141,7 +142,7 @@ public Block copyPositions(int[] positions, int offset, int length)
                 newOffsets,
                 newKeys,
                 newValues,
-                Optional.ofNullable(newHashTable),
+                new HashTables(newHashTable),
                 keyType,
                 keyBlockNativeEquals,
                 keyNativeHashCode,
@@ -161,7 +162,7 @@ public Block getRegion(int position, int length)
                 getOffsets(),
                 getRawKeyBlock(),
                 getRawValueBlock(),
-                Optional.ofNullable(getHashTables()),
+                getHashTables(),
                 keyType,
                 keyBlockNativeEquals,
                 keyNativeHashCode,
@@ -181,7 +182,8 @@ public long getRegionSizeInBytes(int position, int length)
         return getRawKeyBlock().getRegionSizeInBytes(entriesStart, entryCount) +
                 getRawValueBlock().getRegionSizeInBytes(entriesStart, entryCount) +
                 (Integer.BYTES + Byte.BYTES) * (long) length +
-                Integer.BYTES * HASH_MULTIPLIER * (long) entryCount;
+                Integer.BYTES * HASH_MULTIPLIER * (long) entryCount +
+                getHashTables().getInstanceSizeInBytes();
     }
 
     @Override
@@ -211,7 +213,8 @@ public long getPositionsSizeInBytes(boolean[] positions)
         return getRawKeyBlock().getPositionsSizeInBytes(entryPositions) +
                 getRawValueBlock().getPositionsSizeInBytes(entryPositions) +
                 (Integer.BYTES + Byte.BYTES) * (long) usedPositionCount +
-                Integer.BYTES * HASH_MULTIPLIER * (long) usedEntryCount;
+                Integer.BYTES * HASH_MULTIPLIER * (long) usedEntryCount +
+                getHashTables().getInstanceSizeInBytes();
     }
     @Override
     public Block copyRegion(int position, int length)
@@ -228,13 +231,13 @@ public Block copyRegion(int position, int length)
         boolean[] mapIsNull = getMapIsNull();
         boolean[] newMapIsNull = mapIsNull == null ? null : compactArray(mapIsNull, position + getOffsetBase(), length);
 
-        int[] hashTables = getHashTables();
+        int[] hashTables = getHashTables().getRawHashTables();
         int[] newHashTable = null;
         if (hashTables != null) {
             newHashTable = compactArray(hashTables, startValueOffset * HASH_MULTIPLIER, (endValueOffset - startValueOffset) * HASH_MULTIPLIER);
         }
 
-        if (newKeys == getRawKeyBlock() && newValues == getRawValueBlock() && newOffsets == getOffsets() && newMapIsNull == mapIsNull && newHashTable == getHashTables()) {
+        if (newKeys == getRawKeyBlock() && newValues == getRawValueBlock() && newOffsets == getOffsets() && newMapIsNull == mapIsNull && newHashTable == hashTables) {
             return this;
         }
         return createMapBlockInternal(
@@ -244,7 +247,7 @@ public Block copyRegion(int position, int length)
                 newOffsets,
                 newKeys,
                 newValues,
-                Optional.ofNullable(newHashTable),
+                new HashTables(newHashTable),
                 keyType,
                 keyBlockNativeEquals,
                 keyNativeHashCode,
@@ -285,7 +288,7 @@ public Block getSingleValueBlock(int position)
         Block newKeys = getRawKeyBlock().copyRegion(startValueOffset, valueLength);
         Block newValues = getRawValueBlock().copyRegion(startValueOffset, valueLength);
 
-        int[] hashTables = getHashTables();
+        int[] hashTables = getHashTables().getRawHashTables();
         int[] newHashTable = null;
         if (hashTables != null) {
             newHashTable = Arrays.copyOfRange(hashTables, startValueOffset * HASH_MULTIPLIER, endValueOffset * HASH_MULTIPLIER);
@@ -298,7 +301,7 @@ public Block getSingleValueBlock(int position)
                 new int[] {0, valueLength},
                 newKeys,
                 newValues,
-                Optional.ofNullable(newHashTable),
+                new HashTables(newHashTable),
                 keyType,
                 keyBlockNativeEquals,
                 keyNativeHashCode,
@@ -335,10 +338,55 @@ public boolean isNull(int position)
         return mapIsNull != null && mapIsNull[position + getOffsetBase()];
     }
 
+    // Only for testing
+    public boolean isHashTablesCreated()
+    {
+        return getHashTables().getRawHashTables() != null;
+    }
+
     private void checkReadablePosition(int position)
     {
         if (position < 0 || position >= getPositionCount()) {
             throw new IllegalArgumentException("position is not valid");
         }
     }
+
+    public static class HashTables
+    {
+        private static final int INSTANCE_SIZE = ClassLayout.parseClass(MapBlock.class).instanceSize();
+
+        // Hash to location in map. Writes to the field is protected by "this" monitor.
+        @Nullable
+        private volatile int[] hashTables;
+
+        HashTables()
+        {
+            this.hashTables = null;
+        }
+
+        HashTables(int[] hashTables)
+        {
+            this.hashTables = hashTables;
+        }
+
+        int[] getRawHashTables()
+        {
+            return hashTables;
+        }
+
+        void setRawHashTables(int[] hashTables)
+        {
+            this.hashTables = hashTables;
+        }
+
+        public long getInstanceSizeInBytes()
+        {
+            return INSTANCE_SIZE;
+        }
+
+        public long getRetainedSizeInBytes()
+        {
+            return INSTANCE_SIZE + sizeOf(hashTables);
+        }
+    }
 }

presto-spi/src/main/java/com/facebook/presto/spi/block/MapBlock.java

@@ -43,7 +43,7 @@ public class MapBlock
     private final int[] offsets;
     private final Block keyBlock;
     private final Block valueBlock;
-    private volatile int[] hashTables; // hash to location in map. Writes to the field is protected by "this" monitor.
+    private HashTables hashTables;
 
     private volatile long sizeInBytes;
     private final long retainedSizeInBytes;
@@ -77,7 +77,7 @@ public static MapBlock fromKeyValueBlock(
                 offsets,
                 keyBlock,
                 valueBlock,
-                Optional.empty(),
+                new HashTables(),
                 mapType.getKeyType(),
                 keyBlockNativeEquals,
                 keyNativeHashCode,
@@ -100,21 +100,22 @@ public static MapBlock createMapBlockInternal(
             int[] offsets,
             Block keyBlock,
             Block valueBlock,
-            Optional<int[]> hashTables,
+            HashTables hashTables,
             Type keyType,
             MethodHandle keyBlockNativeEquals,
             MethodHandle keyNativeHashCode,
             MethodHandle keyBlockHashCode)
     {
         validateConstructorArguments(startOffset, positionCount, mapIsNull.orElse(null), offsets, keyBlock, valueBlock, keyType, keyBlockNativeEquals, keyNativeHashCode);
+        requireNonNull(hashTables, "hashTables is null");
         return new MapBlock(
                 startOffset,
                 positionCount,
                 mapIsNull.orElse(null),
                 offsets,
                 keyBlock,
                 valueBlock,
-                hashTables.orElse(null),
+                hashTables,
                 keyType,
                 keyBlockNativeEquals,
                 keyNativeHashCode,
@@ -171,16 +172,17 @@ private MapBlock(
             int[] offsets,
             Block keyBlock,
             Block valueBlock,
-            @Nullable int[] hashTables,
+            HashTables hashTables,
             Type keyType,
             MethodHandle keyBlockNativeEquals,
             MethodHandle keyNativeHashCode,
             MethodHandle keyBlockHashCode)
     {
         super(keyType, keyNativeHashCode, keyBlockNativeEquals, keyBlockHashCode);
 
-        if (hashTables != null && hashTables.length < keyBlock.getPositionCount() * HASH_MULTIPLIER) {
-            throw new IllegalArgumentException(format("keyBlock/valueBlock size does not match hash table size: %s %s", keyBlock.getPositionCount(), hashTables.length));
+        int[] rawHashTables = hashTables.getRawHashTables();
+        if (rawHashTables != null && rawHashTables.length < keyBlock.getPositionCount() * HASH_MULTIPLIER) {
+            throw new IllegalArgumentException(format("keyBlock/valueBlock size does not match hash table size: %s %s", keyBlock.getPositionCount(), rawHashTables.length));
         }
 
         this.startOffset = startOffset;
@@ -190,7 +192,6 @@ private MapBlock(
         this.keyBlock = keyBlock;
         this.valueBlock = valueBlock;
         this.hashTables = hashTables;
-
         this.sizeInBytes = -1;
 
         // We will add the hashtable size to the retained size even if it's not built yet. This could be overestimating
@@ -201,7 +202,8 @@ private MapBlock(
                 + valueBlock.getRetainedSizeInBytes()
                 + sizeOf(offsets)
                 + sizeOf(mapIsNull)
-                + sizeOfIntArray(keyBlock.getPositionCount() * HASH_MULTIPLIER);  // hashtable size if it was built
+                + sizeOfIntArray(keyBlock.getPositionCount() * HASH_MULTIPLIER)  // Raw int[] hashTables size if it was built
+                + hashTables.getInstanceSizeInBytes();  // The instance size for HashTables object because it always is not null.
     }
 
     @Override
@@ -217,7 +219,7 @@ protected Block getRawValueBlock()
     }
 
     @Override
-    protected int[] getHashTables()
+    protected HashTables getHashTables()
     {
         return hashTables;
     }
@@ -264,7 +266,8 @@ private void calculateSize()
         sizeInBytes = keyBlock.getRegionSizeInBytes(entriesStart, entryCount) +
                 valueBlock.getRegionSizeInBytes(entriesStart, entryCount) +
                 (Integer.BYTES + Byte.BYTES) * (long) this.positionCount +
-                Integer.BYTES * HASH_MULTIPLIER * (long) entryCount;
+                Integer.BYTES * HASH_MULTIPLIER * (long) entryCount +
+                hashTables.getInstanceSizeInBytes();
     }
 
     @Override
@@ -280,7 +283,7 @@ public void retainedBytesForEachPart(BiConsumer<Object, Long> consumer)
         consumer.accept(valueBlock, valueBlock.getRetainedSizeInBytes());
         consumer.accept(offsets, sizeOf(offsets));
         consumer.accept(mapIsNull, sizeOf(mapIsNull));
-        consumer.accept(hashTables, sizeOf(hashTables));
+        consumer.accept(hashTables, hashTables.getRetainedSizeInBytes());
         consumer.accept(this, (long) INSTANCE_SIZE);
     }
 
@@ -312,7 +315,7 @@ public Block getLoadedBlock()
                 offsets,
                 keyBlock,
                 loadedValueBlock,
-                Optional.ofNullable(hashTables),
+                hashTables,
                 keyType,
                 keyBlockNativeEquals,
                 keyNativeHashCode,
@@ -322,13 +325,13 @@ public Block getLoadedBlock()
     @Override
     protected void ensureHashTableLoaded()
     {
-        if (this.hashTables != null) {
+        if (this.hashTables.getRawHashTables() != null) {
             return;
         }
 
         // This can only happen for MapBlock, not MapBlockBuilder because the latter always has non-null hashtables
         synchronized (this) {
-            if (this.hashTables != null) {
+            if (this.hashTables.getRawHashTables() != null) {
                 return;
             }
 
@@ -352,7 +355,7 @@ protected void ensureHashTableLoaded()
                         keyOffset * HASH_MULTIPLIER,
                         keyCount * HASH_MULTIPLIER);
             }
-            this.hashTables = hashTables;
+            this.hashTables.setRawHashTables(hashTables);
         }
     }
 }