Search in sources :

Example 1 with CatalogLocator

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;
        }
    }
}
Also used : CatalogLocator(com.palantir.foundry.athena.api.CatalogLocator) GetPartitionsResponsePage(com.palantir.foundry.athena.api.GetPartitionsResponsePage) BlockWriter(com.amazonaws.athena.connector.lambda.data.BlockWriter) CatalogLocator(com.palantir.foundry.athena.api.CatalogLocator) FoundryAthenaMetadataServiceBlocking(com.palantir.foundry.athena.api.FoundryAthenaMetadataServiceBlocking) GetTableLayoutRequest(com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest) Optional(java.util.Optional) QueryStatusChecker(com.amazonaws.athena.connector.lambda.QueryStatusChecker) VisibleForTesting(com.google.common.annotations.VisibleForTesting) GetPartitionsRequest(com.palantir.foundry.athena.api.GetPartitionsRequest) GetPartitionsResponsePage(com.palantir.foundry.athena.api.GetPartitionsResponsePage)

Example 2 with CatalogLocator

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);
        }
    }
}
Also used : CatalogLocator(com.palantir.foundry.athena.api.CatalogLocator) GetSplitsResponse(com.amazonaws.athena.connector.lambda.metadata.GetSplitsResponse) Filter(com.palantir.foundry.athena.api.Filter) Logger(org.slf4j.Logger) CatalogLocator(com.palantir.foundry.athena.api.CatalogLocator) Split(com.amazonaws.athena.connector.lambda.domain.Split) LoggerFactory(org.slf4j.LoggerFactory) GetSlicesRequest(com.palantir.foundry.athena.api.GetSlicesRequest) Set(java.util.Set) SpillLocation(com.amazonaws.athena.connector.lambda.domain.spill.SpillLocation) GetSlicesResponse(com.palantir.foundry.athena.api.GetSlicesResponse) Collectors(java.util.stream.Collectors) AndFilter(com.palantir.foundry.athena.api.AndFilter) HashSet(java.util.HashSet) FoundryAthenaMetadataServiceBlocking(com.palantir.foundry.athena.api.FoundryAthenaMetadataServiceBlocking) EncryptionKey(com.amazonaws.athena.connector.lambda.security.EncryptionKey) Optional(java.util.Optional) GetSplitsRequest(com.amazonaws.athena.connector.lambda.metadata.GetSplitsRequest) GetSlicesResponse(com.palantir.foundry.athena.api.GetSlicesResponse) Filter(com.palantir.foundry.athena.api.Filter) AndFilter(com.palantir.foundry.athena.api.AndFilter) GetSplitsResponse(com.amazonaws.athena.connector.lambda.metadata.GetSplitsResponse) Split(com.amazonaws.athena.connector.lambda.domain.Split) HashSet(java.util.HashSet)

Aggregations

CatalogLocator (com.palantir.foundry.athena.api.CatalogLocator)2 FoundryAthenaMetadataServiceBlocking (com.palantir.foundry.athena.api.FoundryAthenaMetadataServiceBlocking)2 Optional (java.util.Optional)2 QueryStatusChecker (com.amazonaws.athena.connector.lambda.QueryStatusChecker)1 BlockWriter (com.amazonaws.athena.connector.lambda.data.BlockWriter)1 Split (com.amazonaws.athena.connector.lambda.domain.Split)1 SpillLocation (com.amazonaws.athena.connector.lambda.domain.spill.SpillLocation)1 GetSplitsRequest (com.amazonaws.athena.connector.lambda.metadata.GetSplitsRequest)1 GetSplitsResponse (com.amazonaws.athena.connector.lambda.metadata.GetSplitsResponse)1 GetTableLayoutRequest (com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest)1 EncryptionKey (com.amazonaws.athena.connector.lambda.security.EncryptionKey)1 VisibleForTesting (com.google.common.annotations.VisibleForTesting)1 AndFilter (com.palantir.foundry.athena.api.AndFilter)1 Filter (com.palantir.foundry.athena.api.Filter)1 GetPartitionsRequest (com.palantir.foundry.athena.api.GetPartitionsRequest)1 GetPartitionsResponsePage (com.palantir.foundry.athena.api.GetPartitionsResponsePage)1 GetSlicesRequest (com.palantir.foundry.athena.api.GetSlicesRequest)1 GetSlicesResponse (com.palantir.foundry.athena.api.GetSlicesResponse)1 HashSet (java.util.HashSet)1 Set (java.util.Set)1