Examples with GetSlicesResponse com.palantir.foundry.athena.api.GetSlicesResponse used on opensource projects

Example 1 with GetSlicesResponse

use of com.palantir.foundry.athena.api.GetSlicesResponse 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);
        }
    }
}