diff --git a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/InMemoryCachingHiveMetastore.java b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/InMemoryCachingHiveMetastore.java
index 3c0e27cec492..46f29a0574b9 100644
--- a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/InMemoryCachingHiveMetastore.java
+++ b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/InMemoryCachingHiveMetastore.java
@@ -37,6 +37,7 @@
 import javax.inject.Inject;
 
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
@@ -720,15 +721,15 @@ private Map<KeyAndContext<HivePartitionName>, Optional<Partition>> loadPartition
         String tableName = hiveTableName.getTableName();
 
         List<PartitionNameWithVersion> partitionsToFetch = new ArrayList<>();
+        Map<String, PartitionNameWithVersion> partitionNameToVersionMap = new HashMap<>();
         for (KeyAndContext<HivePartitionName> partitionNameKey : partitionNamesKey) {
             checkArgument(partitionNameKey.getKey().getHiveTableName().equals(hiveTableName), "Expected table name %s but got %s", hiveTableName, partitionNameKey.getKey().getHiveTableName());
             checkArgument(partitionNameKey.getContext().equals(firstPartitionKey.getContext()), "Expected context %s but got %s", firstPartitionKey.getContext(), partitionNameKey.getContext());
             partitionsToFetch.add(partitionNameKey.getKey().getPartitionNameWithVersion().get());
+            partitionNameToVersionMap.put(partitionNameKey.getKey().getPartitionNameWithVersion().get().getPartitionName(), partitionNameKey.getKey().getPartitionNameWithVersion().get());
         }
 
         ImmutableMap.Builder<KeyAndContext<HivePartitionName>, Optional<Partition>> partitions = ImmutableMap.builder();
-        ImmutableMap<String, PartitionNameWithVersion> partitionNameToVersionMap = partitionsToFetch.stream()
-                .collect(toImmutableMap(PartitionNameWithVersion::getPartitionName, Function.identity()));
         Map<String, Optional<Partition>> partitionsByNames = delegate.getPartitionsByNames(firstPartitionKey.getContext(), databaseName, tableName, partitionsToFetch);
         for (Entry<String, Optional<Partition>> entry : partitionsByNames.entrySet()) {
             partitions.put(getCachingKey(firstPartitionKey.getContext(), HivePartitionName.hivePartitionName(hiveTableName, partitionNameToVersionMap.get(entry.getKey()))), entry.getValue());