Search in sources :

Example 51 with BlockAllocator

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

the class DataLakeGen2MetadataHandlerTest method doGetTableLayoutWithNoPartitions.

@Test
public void doGetTableLayoutWithNoPartitions() throws Exception {
    BlockAllocator blockAllocator = new BlockAllocatorImpl();
    Constraints constraints = Mockito.mock(Constraints.class);
    TableName tableName = new TableName("testSchema", "testTable");
    Schema partitionSchema = this.dataLakeGen2MetadataHandler.getPartitionSchema("testCatalogName");
    Set<String> partitionCols = partitionSchema.getFields().stream().map(Field::getName).collect(Collectors.toSet());
    GetTableLayoutRequest getTableLayoutRequest = new GetTableLayoutRequest(this.federatedIdentity, "testQueryId", "testCatalogName", tableName, constraints, partitionSchema, partitionCols);
    GetTableLayoutResponse getTableLayoutResponse = this.dataLakeGen2MetadataHandler.doGetTableLayout(blockAllocator, getTableLayoutRequest);
    List<String> actualValues = new ArrayList<>();
    for (int i = 0; i < getTableLayoutResponse.getPartitions().getRowCount(); i++) {
        actualValues.add(BlockUtils.rowToString(getTableLayoutResponse.getPartitions(), i));
    }
    Assert.assertEquals(Collections.singletonList("[PARTITION_NUMBER : 0]"), actualValues);
    SchemaBuilder expectedSchemaBuilder = SchemaBuilder.newBuilder();
    expectedSchemaBuilder.addField(FieldBuilder.newBuilder(DataLakeGen2MetadataHandler.PARTITION_NUMBER, org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build());
    Schema expectedSchema = expectedSchemaBuilder.build();
    Assert.assertEquals(expectedSchema, getTableLayoutResponse.getPartitions().getSchema());
    Assert.assertEquals(tableName, getTableLayoutResponse.getTableName());
}
Also used : Schema(org.apache.arrow.vector.types.pojo.Schema) ArrayList(java.util.ArrayList) TableName(com.amazonaws.athena.connector.lambda.domain.TableName) Constraints(com.amazonaws.athena.connector.lambda.domain.predicate.Constraints) GetTableLayoutResponse(com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutResponse) BlockAllocatorImpl(com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl) BlockAllocator(com.amazonaws.athena.connector.lambda.data.BlockAllocator) GetTableLayoutRequest(com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest) SchemaBuilder(com.amazonaws.athena.connector.lambda.data.SchemaBuilder) Test(org.junit.Test)

Aggregations

BlockAllocator (com.amazonaws.athena.connector.lambda.data.BlockAllocator)51 BlockAllocatorImpl (com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl)50 Test (org.junit.Test)43 Schema (org.apache.arrow.vector.types.pojo.Schema)39 TableName (com.amazonaws.athena.connector.lambda.domain.TableName)38 ResultSet (java.sql.ResultSet)36 AtomicInteger (java.util.concurrent.atomic.AtomicInteger)36 Constraints (com.amazonaws.athena.connector.lambda.domain.predicate.Constraints)35 PreparedStatement (java.sql.PreparedStatement)27 GetTableLayoutRequest (com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest)26 GetTableLayoutResponse (com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutResponse)25 SchemaBuilder (com.amazonaws.athena.connector.lambda.data.SchemaBuilder)20 Map (java.util.Map)16 GetSplitsRequest (com.amazonaws.athena.connector.lambda.metadata.GetSplitsRequest)15 GetSplitsResponse (com.amazonaws.athena.connector.lambda.metadata.GetSplitsResponse)15 HashSet (java.util.HashSet)15 ArrayList (java.util.ArrayList)13 ObjectMapper (com.fasterxml.jackson.databind.ObjectMapper)10 FederationRequest (com.amazonaws.athena.connector.lambda.request.FederationRequest)6 IOException (java.io.IOException)5