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());
}
Aggregations