Examples with MAX_SPLITS_PER_REQUEST com.amazonaws.athena.connectors.dynamodb.DynamoDBMetadataHandler.MAX_SPLITS_PER_REQUEST used on opensource projects

Example 1 with MAX_SPLITS_PER_REQUEST

use of com.amazonaws.athena.connectors.dynamodb.DynamoDBMetadataHandler.MAX_SPLITS_PER_REQUEST in project aws-athena-query-federation by awslabs.

the class DynamoDBMetadataHandlerTest method doGetSplitsQuery.

@Test
public void doGetSplitsQuery() throws Exception {
    Map<String, ValueSet> constraintsMap = new HashMap<>();
    EquatableValueSet.Builder valueSet = EquatableValueSet.newBuilder(allocator, Types.MinorType.VARCHAR.getType(), true, false);
    for (int i = 0; i < 2000; i++) {
        valueSet.add("test_str_" + i);
    }
    constraintsMap.put("col_0", valueSet.build());
    GetTableLayoutResponse layoutResponse = handler.doGetTableLayout(allocator, new GetTableLayoutRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, TEST_TABLE_NAME, new Constraints(constraintsMap), SchemaBuilder.newBuilder().build(), Collections.EMPTY_SET));
    GetSplitsRequest req = new GetSplitsRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, TEST_TABLE_NAME, layoutResponse.getPartitions(), ImmutableList.of("col_0"), new Constraints(new HashMap<>()), null);
    logger.info("doGetSplits: req[{}]", req);
    GetSplitsResponse response = handler.doGetSplits(allocator, req);
    assertThat(response.getRequestType(), equalTo(MetadataRequestType.GET_SPLITS));
    String continuationToken = response.getContinuationToken();
    logger.info("doGetSplits: continuationToken[{}] - numSplits[{}]", continuationToken, response.getSplits().size());
    assertThat(continuationToken, equalTo(String.valueOf(MAX_SPLITS_PER_REQUEST - 1)));
    assertThat(response.getSplits().size(), equalTo(MAX_SPLITS_PER_REQUEST));
    assertThat(response.getSplits().stream().map(split -> split.getProperty("col_0")).distinct().count(), equalTo((long) MAX_SPLITS_PER_REQUEST));
    response = handler.doGetSplits(allocator, new GetSplitsRequest(req, continuationToken));
    logger.info("doGetSplits: continuationToken[{}] - numSplits[{}]", continuationToken, response.getSplits().size());
    assertThat(response.getContinuationToken(), equalTo(null));
    assertThat(response.getSplits().size(), equalTo(MAX_SPLITS_PER_REQUEST));
    assertThat(response.getSplits().stream().map(split -> split.getProperty("col_0")).distinct().count(), equalTo((long) MAX_SPLITS_PER_REQUEST));
}