Examples with BlockSpiller com.amazonaws.athena.connector.lambda.data.BlockSpiller used on opensource projects

Example 11 with BlockSpiller

use of com.amazonaws.athena.connector.lambda.data.BlockSpiller in project aws-athena-query-federation by awslabs.

the class AbstractTableProviderTest method readTableTest.

@Test
public void readTableTest() {
    GetTableRequest request = new GetTableRequest(identity, expectedQuery, expectedCatalog, expectedTableName);
    GetTableResponse response = provider.getTable(allocator, request);
    assertTrue(response.getSchema().getFields().size() > 1);
    Map<String, ValueSet> constraintsMap = new HashMap<>();
    constraintsMap.put(idField, EquatableValueSet.newBuilder(allocator, Types.MinorType.VARCHAR.getType(), true, false).add(idValue).build());
    Constraints constraints = new Constraints(constraintsMap);
    ConstraintEvaluator evaluator = new ConstraintEvaluator(allocator, response.getSchema(), constraints);
    S3SpillLocation spillLocation = S3SpillLocation.newBuilder().withBucket("bucket").withPrefix("prefix").withSplitId(UUID.randomUUID().toString()).withQueryId(UUID.randomUUID().toString()).withIsDirectory(true).build();
    ReadRecordsRequest readRequest = new ReadRecordsRequest(identity, expectedCatalog, "queryId", expectedTableName, response.getSchema(), Split.newBuilder(spillLocation, keyFactory.create()).build(), constraints, 100_000_000, 100_000_000);
    SpillConfig spillConfig = SpillConfig.newBuilder().withSpillLocation(spillLocation).withMaxBlockBytes(3_000_000).withMaxInlineBlockBytes(0).withRequestId("queryid").withEncryptionKey(keyFactory.create()).build();
    setUpRead();
    BlockSpiller spiller = new S3BlockSpiller(amazonS3, spillConfig, allocator, response.getSchema(), evaluator);
    provider.readWithConstraint(spiller, readRequest, queryStatusChecker);
    validateRead(response.getSchema(), blockSpillReader, spiller.getSpillLocations(), spillConfig.getEncryptionKey());
}

Example 12 with BlockSpiller

use of com.amazonaws.athena.connector.lambda.data.BlockSpiller in project aws-athena-query-federation by awslabs.

the class HiveMuxRecordHandlerTest method readWithConstraintWithUnsupportedCatalog.

@Test(expected = RuntimeException.class)
public void readWithConstraintWithUnsupportedCatalog() throws SQLException {
    BlockSpiller blockSpiller = Mockito.mock(BlockSpiller.class);
    ReadRecordsRequest readRecordsRequest = Mockito.mock(ReadRecordsRequest.class);
    Mockito.when(readRecordsRequest.getCatalogName()).thenReturn("unsupportedCatalog");
    this.jdbcRecordHandler.readWithConstraint(blockSpiller, readRecordsRequest, queryStatusChecker);
}

Example 13 with BlockSpiller

use of com.amazonaws.athena.connector.lambda.data.BlockSpiller in project aws-athena-query-federation by awslabs.

the class ImpalaMuxRecordHandlerTest method readWithConstraintWithUnsupportedCatalog.

@Test(expected = RuntimeException.class)
public void readWithConstraintWithUnsupportedCatalog() throws SQLException {
    BlockSpiller blockSpiller = Mockito.mock(BlockSpiller.class);
    ReadRecordsRequest readRecordsRequest = Mockito.mock(ReadRecordsRequest.class);
    Mockito.when(readRecordsRequest.getCatalogName()).thenReturn("unsupportedCatalog");
    this.jdbcRecordHandler.readWithConstraint(blockSpiller, readRecordsRequest, queryStatusChecker);
}

Example 14 with BlockSpiller

use of com.amazonaws.athena.connector.lambda.data.BlockSpiller in project aws-athena-query-federation by awslabs.

the class ImpalaMuxRecordHandlerTest method readWithConstraint.

@Test
public void readWithConstraint() throws SQLException {
    BlockSpiller blockSpiller = Mockito.mock(BlockSpiller.class);
    ReadRecordsRequest readRecordsRequest = Mockito.mock(ReadRecordsRequest.class);
    Mockito.when(readRecordsRequest.getCatalogName()).thenReturn("recordImpala");
    this.jdbcRecordHandler.readWithConstraint(blockSpiller, readRecordsRequest, queryStatusChecker);
    Mockito.verify(this.impalaRecordHandler, Mockito.times(1)).readWithConstraint(Mockito.eq(blockSpiller), Mockito.eq(readRecordsRequest), Mockito.eq(queryStatusChecker));
}

Example 15 with BlockSpiller

use of com.amazonaws.athena.connector.lambda.data.BlockSpiller in project aws-athena-query-federation by awslabs.

the class DataLakeGen2MuxRecordHandlerTest method readWithConstraint.

@Test
public void readWithConstraint() {
    BlockSpiller blockSpiller = Mockito.mock(BlockSpiller.class);
    ReadRecordsRequest readRecordsRequest = Mockito.mock(ReadRecordsRequest.class);
    Mockito.when(readRecordsRequest.getCatalogName()).thenReturn(DataLakeGen2Constants.NAME);
    this.jdbcRecordHandler.readWithConstraint(blockSpiller, readRecordsRequest, queryStatusChecker);
    Mockito.verify(this.dataLakeGen2RecordHandler, Mockito.times(1)).readWithConstraint(Mockito.eq(blockSpiller), Mockito.eq(readRecordsRequest), Mockito.eq(queryStatusChecker));
}