diff --git a/presto-hive/src/main/java/io/prestosql/plugin/hive/metastore/glue/GlueHiveMetastore.java b/presto-hive/src/main/java/io/prestosql/plugin/hive/metastore/glue/GlueHiveMetastore.java index e54e7d0251cd..ff9e0dd9ec59 100644 --- a/presto-hive/src/main/java/io/prestosql/plugin/hive/metastore/glue/GlueHiveMetastore.java +++ b/presto-hive/src/main/java/io/prestosql/plugin/hive/metastore/glue/GlueHiveMetastore.java @@ -870,24 +870,39 @@ private Map> getPartitionsByNames(Table table, List< private List batchGetPartition(Table table, List partitionNames) { try { - List> batchGetPartitionFutures = new ArrayList<>(); - - for (List partitionNamesBatch : Lists.partition(partitionNames, BATCH_GET_PARTITION_MAX_PAGE_SIZE)) { - List partitionValuesBatch = mappedCopy(partitionNamesBatch, partitionName -> new PartitionValueList().withValues(toPartitionValues(partitionName))); - batchGetPartitionFutures.add(glueClient.batchGetPartitionAsync(new BatchGetPartitionRequest() - .withCatalogId(catalogId) - .withDatabaseName(table.getDatabaseName()) - .withTableName(table.getTableName()) - .withPartitionsToGet(partitionValuesBatch))); - } + List partitionValueLists = mappedCopy(partitionNames, partitionName -> new PartitionValueList().withValues(toPartitionValues(partitionName))); + ImmutableList.Builder resultsBuilder = ImmutableList.builderWithExpectedSize(partitionNames.size()); // Reuse immutable field instances opportunistically between partitions GluePartitionConverter converter = new GluePartitionConverter(table); - ImmutableList.Builder resultsBuilder = ImmutableList.builderWithExpectedSize(partitionNames.size()); - for (Future future : batchGetPartitionFutures) { - future.get().getPartitions().stream() - .map(converter) - .forEach(resultsBuilder::add); + + while (!partitionValueLists.isEmpty()) { + List> batchedPartitionValueLists = Lists.partition(partitionValueLists, BATCH_GET_PARTITION_MAX_PAGE_SIZE); + List> batchGetPartitionFutures = new ArrayList<>(); + for (List partitions : batchedPartitionValueLists) { + batchGetPartitionFutures.add(glueClient.batchGetPartitionAsync(new BatchGetPartitionRequest() + .withCatalogId(catalogId) + .withDatabaseName(table.getDatabaseName()) + .withTableName(table.getTableName()) + .withPartitionsToGet(partitions))); + } + + List unprocessedPartitions = new ArrayList<>(); + for (Future future : batchGetPartitionFutures) { + BatchGetPartitionResult batchGetPartitionResult = future.get(); + + // In the unlikely scenario where batchGetPartition call cannot make progress on retrieving partitions, + // we throw exception on this metastore operation to avoid getting stuck in this loop. + if (batchGetPartitionResult.getPartitions().isEmpty() && !batchGetPartitionResult.getUnprocessedKeys().isEmpty()) { + throw new PrestoException(HIVE_METASTORE_ERROR, "Unable to retrieve partitions: " + batchGetPartitionResult.getUnprocessedKeys()); + } + + batchGetPartitionResult.getPartitions().stream() + .map(converter) + .forEach(resultsBuilder::add); + unprocessedPartitions.addAll(batchGetPartitionResult.getUnprocessedKeys()); + } + partitionValueLists = unprocessedPartitions; } return resultsBuilder.build();