use of com.palantir.foundry.athena.api.CatalogLocator in project foundry-athena-query-federation-connector by palantir.
the class PartitionFetcher method getAndWritePartitions.
private void getAndWritePartitions(BlockWriter blockWriter, GetTableLayoutRequest request, QueryStatusChecker queryStatusChecker) {
CatalogLocator locator = FoundryAthenaObjectMapper.objectMapper().convertValue(request.getSchema().getCustomMetadata(), CatalogLocator.class);
Optional<String> pageToken = Optional.empty();
while (queryStatusChecker.isQueryRunning()) {
GetPartitionsResponsePage page = metadataService.getPartitions(authProvider.getAuthHeader(), GetPartitionsRequest.builder().locator(locator).limit(PAGE_SIZE).pageToken(pageToken).build());
page.getPartitions().forEach(partition -> blockWriter.writeRows((block, rowNum) -> {
boolean matched = partition.get().entrySet().stream().map(fieldName -> fieldName.getValue().accept(new PartitionValueWriter(block, fieldName.getKey(), rowNum))).reduce(true, Boolean::logicalAnd);
// if all fields passed then we wrote 1 row
return matched ? 1 : 0;
}));
if (page.getNextPageToken().isPresent()) {
pageToken = page.getNextPageToken();
} else {
return;
}
}
}
use of com.palantir.foundry.athena.api.CatalogLocator in project foundry-athena-query-federation-connector by palantir.
the class SplitsFetcher method getSplits.
GetSplitsResponse getSplits(GetSplitsRequest request, SpillLocationFactory spillLocationFactory, EncryptionKey encryptionKey) {
CatalogLocator locator = FoundryAthenaObjectMapper.objectMapper().convertValue(request.getSchema().getCustomMetadata(), CatalogLocator.class);
Optional<Filter> filter;
if (request.getConstraints().getSummary().isEmpty()) {
filter = Optional.empty();
} else {
// we just push down all constraints which will include those for any partition columns
filter = Optional.of(Filter.and(AndFilter.of(request.getConstraints().getSummary().entrySet().stream().map(entry -> ConstraintConverter.convert(entry.getKey(), entry.getValue())).collect(Collectors.toList()))));
}
Set<Split> splits = new HashSet<>();
Optional<String> pageToken = Optional.empty();
while (true) {
GetSlicesResponse response = metadataService.getSlices(authProvider.getAuthHeader(), GetSlicesRequest.builder().locator(locator).filter(filter).nextPageToken(pageToken).build());
splits.addAll(response.getSlices().stream().map(slice -> slices.toSplit(spillLocationFactory.makeSpillLocation(), encryptionKey, slice)).collect(Collectors.toSet()));
if (response.getNextPageToken().isPresent()) {
pageToken = response.getNextPageToken();
} else {
log.debug("finished planning splits. number of splits: {}", splits.size());
return new GetSplitsResponse(request.getCatalogName(), splits);
}
}
}
Aggregations