diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueHiveMetastore.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueHiveMetastore.java index 0b6003180bcd..5682814aeec1 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueHiveMetastore.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueHiveMetastore.java @@ -128,6 +128,7 @@ import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Strings.isNullOrEmpty; +import static com.google.common.base.Verify.verify; import static com.google.common.collect.Comparators.lexicographical; import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.ImmutableMap.toImmutableMap; @@ -149,6 +150,7 @@ import static java.util.Comparator.comparing; import static java.util.Objects.requireNonNull; import static java.util.function.UnaryOperator.identity; +import static java.util.stream.Collectors.toCollection; import static java.util.stream.Collectors.toMap; import static org.apache.hadoop.hive.metastore.TableType.MANAGED_TABLE; import static org.apache.hadoop.hive.metastore.TableType.VIRTUAL_VIEW; @@ -907,24 +909,42 @@ 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 pendingPartitions = partitionNames.stream() + .map(partitionName -> new PartitionValueList().withValues(toPartitionValues(partitionName))) + .collect(toCollection(ArrayList::new)); + + 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 (!pendingPartitions.isEmpty()) { + List> batchGetPartitionFutures = new ArrayList<>(); + for (List partitions : Lists.partition(pendingPartitions, BATCH_GET_PARTITION_MAX_PAGE_SIZE)) { + batchGetPartitionFutures.add(glueClient.batchGetPartitionAsync(new BatchGetPartitionRequest() + .withCatalogId(catalogId) + .withDatabaseName(table.getDatabaseName()) + .withTableName(table.getTableName()) + .withPartitionsToGet(partitions))); + } + pendingPartitions.clear(); + + for (Future future : batchGetPartitionFutures) { + BatchGetPartitionResult batchGetPartitionResult = future.get(); + List partitions = batchGetPartitionResult.getPartitions(); + List unprocessedKeys = batchGetPartitionResult.getUnprocessedKeys(); + + // In the unlikely scenario where batchGetPartition call cannot make progress on retrieving partitions, avoid infinite loop + if (partitions.isEmpty()) { + verify(!unprocessedKeys.isEmpty(), "Empty unprocessedKeys for non-empty BatchGetPartitionRequest and empty partitions result"); + throw new TrinoException(HIVE_METASTORE_ERROR, "Cannot make progress retrieving partitions. Unable to retrieve partitions: " + unprocessedKeys); + } + + partitions.stream() + .map(converter) + .forEach(resultsBuilder::add); + pendingPartitions.addAll(unprocessedKeys); + } } return resultsBuilder.build();