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

Example 21 with BlockAllocator

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

the class SqlServerMetadataHandlerTest method doGetTable.

@Test
public void doGetTable() throws Exception {
    BlockAllocator blockAllocator = new BlockAllocatorImpl();
    String[] schema = { "DATA_TYPE", "COLUMN_SIZE", "COLUMN_NAME", "DECIMAL_DIGITS", "NUM_PREC_RADIX" };
    Object[][] values = { { Types.INTEGER, 12, "testCol1", 0, 0 }, { Types.VARCHAR, 25, "testCol2", 0, 0 }, { Types.TIMESTAMP, 93, "testCol3", 0, 0 }, { Types.TIMESTAMP_WITH_TIMEZONE, 93, "testCol4", 0, 0 } };
    AtomicInteger rowNumber = new AtomicInteger(-1);
    ResultSet resultSet = mockResultSet(schema, values, rowNumber);
    SchemaBuilder expectedSchemaBuilder = SchemaBuilder.newBuilder();
    expectedSchemaBuilder.addField(FieldBuilder.newBuilder("testCol1", org.apache.arrow.vector.types.Types.MinorType.INT.getType()).build());
    expectedSchemaBuilder.addField(FieldBuilder.newBuilder("testCol2", org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build());
    expectedSchemaBuilder.addField(FieldBuilder.newBuilder("testCol3", org.apache.arrow.vector.types.Types.MinorType.DATEMILLI.getType()).build());
    expectedSchemaBuilder.addField(FieldBuilder.newBuilder("testCol4", org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build());
    PARTITION_SCHEMA.getFields().forEach(expectedSchemaBuilder::addField);
    Schema expected = expectedSchemaBuilder.build();
    TableName inputTableName = new TableName("TESTSCHEMA", "TESTTABLE");
    Mockito.when(connection.getMetaData().getColumns("testCatalog", inputTableName.getSchemaName(), inputTableName.getTableName(), null)).thenReturn(resultSet);
    Mockito.when(connection.getCatalog()).thenReturn("testCatalog");
    GetTableResponse getTableResponse = this.sqlServerMetadataHandler.doGetTable(blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName));
    Assert.assertEquals(expected, getTableResponse.getSchema());
    Assert.assertEquals(inputTableName, getTableResponse.getTableName());
    Assert.assertEquals("testCatalog", getTableResponse.getCatalogName());
}

Example 22 with BlockAllocator

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

the class SqlServerMetadataHandlerTest 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.sqlServerMetadataHandler.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);
    PreparedStatement preparedStatement = Mockito.mock(PreparedStatement.class);
    Mockito.when(this.connection.prepareStatement(sqlServerMetadataHandler.GET_PARTITIONS_QUERY)).thenReturn(preparedStatement);
    String[] columns = { sqlServerMetadataHandler.PARTITION_NUMBER };
    int[] types = { Types.VARCHAR };
    Object[][] values = { {} };
    ResultSet resultSet = mockResultSet(columns, types, values, new AtomicInteger(-1));
    Mockito.when(preparedStatement.executeQuery()).thenReturn(resultSet);
    PreparedStatement rowCountPreparedStatement = Mockito.mock(PreparedStatement.class);
    Mockito.when(this.connection.prepareStatement(sqlServerMetadataHandler.ROW_COUNT_QUERY)).thenReturn(rowCountPreparedStatement);
    ResultSet rowCountResultSet = mockResultSet(new String[] { "ROW_COUNT" }, new int[] { Types.INTEGER }, new Object[][] { { 0 } }, new AtomicInteger(-1));
    Mockito.when(rowCountPreparedStatement.executeQuery()).thenReturn(rowCountResultSet);
    Mockito.when(this.connection.getMetaData().getSearchStringEscape()).thenReturn(null);
    GetTableLayoutResponse getTableLayoutResponse = this.sqlServerMetadataHandler.doGetTableLayout(blockAllocator, getTableLayoutRequest);
    Assert.assertEquals(values.length, getTableLayoutResponse.getPartitions().getRowCount());
    List<String> actualValues = new ArrayList<>();
    for (int i = 0; i < getTableLayoutResponse.getPartitions().getRowCount(); i++) {
        actualValues.add(BlockUtils.rowToString(getTableLayoutResponse.getPartitions(), i));
    }
    Assert.assertEquals(Arrays.asList("[PARTITION_NUMBER : 0]"), actualValues);
    SchemaBuilder expectedSchemaBuilder = SchemaBuilder.newBuilder();
    expectedSchemaBuilder.addField(FieldBuilder.newBuilder(sqlServerMetadataHandler.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());
    Mockito.verify(preparedStatement, Mockito.times(1)).setString(1, getTableLayoutRequest.getTableName().getSchemaName() + "." + getTableLayoutRequest.getTableName().getTableName());
}

Example 23 with BlockAllocator

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

the class SqlServerMetadataHandlerTest method doGetSplitsContinuation.

@Test
public void doGetSplitsContinuation() throws Exception {
    BlockAllocator blockAllocator = new BlockAllocatorImpl();
    Constraints constraints = Mockito.mock(Constraints.class);
    TableName tableName = new TableName("testSchema", "testTable");
    Schema partitionSchema = this.sqlServerMetadataHandler.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);
    PreparedStatement rowCountPreparedStatement = Mockito.mock(PreparedStatement.class);
    Mockito.when(this.connection.prepareStatement(sqlServerMetadataHandler.ROW_COUNT_QUERY)).thenReturn(rowCountPreparedStatement);
    ResultSet rowCountResultSet = mockResultSet(new String[] { "ROW_COUNT" }, new int[] { Types.INTEGER }, new Object[][] { { 2 } }, new AtomicInteger(-1));
    Mockito.when(rowCountPreparedStatement.executeQuery()).thenReturn(rowCountResultSet);
    PreparedStatement preparedStatement = Mockito.mock(PreparedStatement.class);
    Mockito.when(this.connection.prepareStatement(sqlServerMetadataHandler.GET_PARTITIONS_QUERY)).thenReturn(preparedStatement);
    String[] columns = { sqlServerMetadataHandler.PARTITION_NUMBER };
    int[] types = { Types.INTEGER };
    Object[][] values = { { 2 }, { 3 } };
    ResultSet resultSet = mockResultSet(columns, types, values, new AtomicInteger(-1));
    Mockito.when(preparedStatement.executeQuery()).thenReturn(resultSet);
    PreparedStatement partFuncPreparedStatement = Mockito.mock(PreparedStatement.class);
    Mockito.when(this.connection.prepareStatement(sqlServerMetadataHandler.GET_PARTITION_FUNCTION_QUERY)).thenReturn(partFuncPreparedStatement);
    ResultSet partFuncResultSet = mockResultSet(new String[] { "PARTITION FUNCTION", "PARTITIONING COLUMN" }, new int[] { Types.VARCHAR, Types.VARCHAR }, new Object[][] { { "pf", "pc" } }, new AtomicInteger(-1));
    Mockito.when(partFuncPreparedStatement.executeQuery()).thenReturn(partFuncResultSet);
    Mockito.when(this.connection.getMetaData().getSearchStringEscape()).thenReturn(null);
    GetTableLayoutResponse getTableLayoutResponse = this.sqlServerMetadataHandler.doGetTableLayout(blockAllocator, getTableLayoutRequest);
    BlockAllocator splitBlockAllocator = new BlockAllocatorImpl();
    GetSplitsRequest getSplitsRequest = new GetSplitsRequest(this.federatedIdentity, "testQueryId", "testCatalogName", tableName, getTableLayoutResponse.getPartitions(), new ArrayList<>(partitionCols), constraints, "2");
    GetSplitsResponse getSplitsResponse = this.sqlServerMetadataHandler.doGetSplits(splitBlockAllocator, getSplitsRequest);
    Set<Map<String, String>> expectedSplits = new HashSet<>();
    expectedSplits.add(Map.ofEntries(Map.entry(sqlServerMetadataHandler.PARTITION_NUMBER, "3"), Map.entry("PARTITIONING_COLUMN", "pc"), Map.entry("PARTITION_FUNCTION", "pf")));
    Set<Map<String, String>> actualSplits = getSplitsResponse.getSplits().stream().map(Split::getProperties).collect(Collectors.toSet());
    Assert.assertEquals(expectedSplits, actualSplits);
}

Example 24 with BlockAllocator

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

the class DataLakeGen2MetadataHandlerTest method doGetSplitsWithNoPartition.

@Test
public void doGetSplitsWithNoPartition() 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);
    BlockAllocator splitBlockAllocator = new BlockAllocatorImpl();
    GetSplitsRequest getSplitsRequest = new GetSplitsRequest(this.federatedIdentity, "testQueryId", "testCatalogName", tableName, getTableLayoutResponse.getPartitions(), new ArrayList<>(partitionCols), constraints, null);
    GetSplitsResponse getSplitsResponse = this.dataLakeGen2MetadataHandler.doGetSplits(splitBlockAllocator, getSplitsRequest);
    Set<Map<String, String>> expectedSplits = new HashSet<>();
    expectedSplits.add(Collections.singletonMap(DataLakeGen2MetadataHandler.PARTITION_NUMBER, "0"));
    Assert.assertEquals(expectedSplits.size(), getSplitsResponse.getSplits().size());
    Set<Map<String, String>> actualSplits = getSplitsResponse.getSplits().stream().map(Split::getProperties).collect(Collectors.toSet());
    Assert.assertEquals(expectedSplits, actualSplits);
}

Example 25 with BlockAllocator

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

the class DataLakeGen2MetadataHandlerTest method doGetTable.

@Test
public void doGetTable() throws Exception {
    BlockAllocator blockAllocator = new BlockAllocatorImpl();
    String[] schema = { "DATA_TYPE", "COLUMN_SIZE", "COLUMN_NAME", "DECIMAL_DIGITS", "NUM_PREC_RADIX" };
    Object[][] values = { { Types.INTEGER, 12, "testCol1", 0, 0 }, { Types.VARCHAR, 25, "testCol2", 0, 0 }, { Types.TIMESTAMP, 93, "testCol3", 0, 0 }, { Types.TIMESTAMP_WITH_TIMEZONE, 93, "testCol4", 0, 0 } };
    AtomicInteger rowNumber = new AtomicInteger(-1);
    ResultSet resultSet = mockResultSet(schema, values, rowNumber);
    SchemaBuilder expectedSchemaBuilder = SchemaBuilder.newBuilder();
    expectedSchemaBuilder.addField(FieldBuilder.newBuilder("testCol1", org.apache.arrow.vector.types.Types.MinorType.INT.getType()).build());
    expectedSchemaBuilder.addField(FieldBuilder.newBuilder("testCol2", org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build());
    expectedSchemaBuilder.addField(FieldBuilder.newBuilder("testCol3", org.apache.arrow.vector.types.Types.MinorType.DATEMILLI.getType()).build());
    expectedSchemaBuilder.addField(FieldBuilder.newBuilder("testCol4", org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build());
    PARTITION_SCHEMA.getFields().forEach(expectedSchemaBuilder::addField);
    Schema expected = expectedSchemaBuilder.build();
    TableName inputTableName = new TableName("TESTSCHEMA", "TESTTABLE");
    Mockito.when(connection.getMetaData().getColumns("testCatalog", inputTableName.getSchemaName(), inputTableName.getTableName(), null)).thenReturn(resultSet);
    Mockito.when(connection.getCatalog()).thenReturn("testCatalog");
    GetTableResponse getTableResponse = this.dataLakeGen2MetadataHandler.doGetTable(blockAllocator, new GetTableRequest(this.federatedIdentity, "testQueryId", "testCatalog", inputTableName));
    Assert.assertEquals(expected, getTableResponse.getSchema());
    Assert.assertEquals(inputTableName, getTableResponse.getTableName());
    Assert.assertEquals("testCatalog", getTableResponse.getCatalogName());
}