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

Example 61 with SchemaBuilder

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

the class TPCDSMetadataHandler method doGetTable.

/**
 * Used to get definition (field names, types, descriptions, etc...) of a Table using the static
 * metadata provided by TerraData's TPCDS generator.
 *
 * @see MetadataHandler
 */
@Override
public GetTableResponse doGetTable(BlockAllocator allocator, GetTableRequest request) {
    logger.info("doGetTable: enter - " + request);
    Table table = TPCDSUtils.validateTable(request.getTableName());
    SchemaBuilder schemaBuilder = SchemaBuilder.newBuilder();
    for (Column nextCol : table.getColumns()) {
        schemaBuilder.addField(TPCDSUtils.convertColumn(nextCol));
    }
    return new GetTableResponse(request.getCatalogName(), request.getTableName(), schemaBuilder.build(), Collections.EMPTY_SET);
}

Example 62 with SchemaBuilder

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

the class TPCDSRecordHandlerTest method doReadRecordForTPCDSTIMETypeColumn.

@Test
public void doReadRecordForTPCDSTIMETypeColumn() throws Exception {
    for (Table next : Table.getBaseTables()) {
        if (next.getName().equals("dbgen_version")) {
            table = next;
        }
    }
    SchemaBuilder schemaBuilder = SchemaBuilder.newBuilder();
    for (Column nextCol : table.getColumns()) {
        schemaBuilder.addField(TPCDSUtils.convertColumn(nextCol));
    }
    ReadRecordsRequest request = new ReadRecordsRequest(identity, "catalog", "queryId-" + System.currentTimeMillis(), new TableName("tpcds1", table.getName()), schemaBuilder.build(), Split.newBuilder(S3SpillLocation.newBuilder().withBucket(UUID.randomUUID().toString()).withSplitId(UUID.randomUUID().toString()).withQueryId(UUID.randomUUID().toString()).withIsDirectory(true).build(), keyFactory.create()).add(SPLIT_NUMBER_FIELD, "0").add(SPLIT_TOTAL_NUMBER_FIELD, "1000").add(SPLIT_SCALE_FACTOR_FIELD, "1").build(), new Constraints(ImmutableMap.of()), 100_000_000_000L, 100_000_000_000L);
    RecordResponse rawResponse = handler.doReadRecords(allocator, request);
    assertTrue(rawResponse instanceof ReadRecordsResponse);
    ReadRecordsResponse response = (ReadRecordsResponse) rawResponse;
    logger.info("doReadRecordForTPCDSTIMETypeColumn: {}", BlockUtils.rowToString(response.getRecords(), 0));
    // TPCDS for `dbgen_version` always generates 1 record.
    assertEquals(1, response.getRecords().getRowCount());
    logger.info("doReadRecordForTPCDSTIMETypeColumn: exit");
}

Example 63 with SchemaBuilder

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

the class VerticaSchemaUtils method buildTableSchema.

// Builds the table schema
protected Schema buildTableSchema(Connection connection, TableName name) {
    try {
        logger.info("Building the schema for table {} ", name);
        SchemaBuilder tableSchemaBuilder = SchemaBuilder.newBuilder();
        DatabaseMetaData dbMetadata = connection.getMetaData();
        ResultSet definition = dbMetadata.getColumns(null, name.getSchemaName(), name.getTableName(), null);
        while (definition.next()) {
            String colType = definition.getString("TYPE_NAME").toUpperCase();
            switch(colType) {
                // If Bit
                case "BIT":
                    {
                        tableSchemaBuilder.addBitField(definition.getString("COLUMN_NAME"));
                        break;
                    }
                // If TinyInt
                case "TINYINT":
                    {
                        tableSchemaBuilder.addTinyIntField(definition.getString("COLUMN_NAME"));
                        break;
                    }
                // If SmallInt
                case "SMALLINT":
                    {
                        tableSchemaBuilder.addSmallIntField(definition.getString("COLUMN_NAME"));
                        break;
                    }
                // If Int
                case "INTEGER":
                // If BIGINT
                case "BIGINT":
                    {
                        tableSchemaBuilder.addBigIntField(definition.getString("COLUMN_NAME"));
                        break;
                    }
                // If FLOAT4
                case "FLOAT4":
                    {
                        tableSchemaBuilder.addFloat4Field(definition.getString("COLUMN_NAME"));
                        break;
                    }
                // If FLOAT8
                case "FLOAT8":
                    {
                        tableSchemaBuilder.addFloat8Field(definition.getString("COLUMN_NAME"));
                        break;
                    }
                // If DECIMAL/NUMERIC
                case "NUMERIC":
                    {
                        tableSchemaBuilder.addDecimalField(definition.getString("COLUMN_NAME"), 10, 2);
                        break;
                    }
                // If VARCHAR
                case "BOOLEAN":
                case "VARCHAR":
                case "TIMESTAMPTZ":
                case "TIMESTAMP":
                    {
                        tableSchemaBuilder.addStringField(definition.getString("COLUMN_NAME"));
                        break;
                    }
                // If DATETIME
                case "DATETIME":
                    {
                        tableSchemaBuilder.addDateDayField(definition.getString("COLUMN_NAME"));
                        break;
                    }
                default:
                    tableSchemaBuilder.addStringField(definition.getString("COLUMN_NAME"));
            }
        }
        return tableSchemaBuilder.build();
    } catch (SQLException e) {
        throw new RuntimeException("Error in building the table schema: " + e.getMessage(), e);
    }
}

Example 64 with SchemaBuilder

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

the class HiveRecordHandlerTest method buildSplitSql.

@Test
public void buildSplitSql() throws SQLException {
    TableName tableName = new TableName("testSchema", "testTable");
    SchemaBuilder schemaBuilder = SchemaBuilder.newBuilder();
    schemaBuilder.addField(FieldBuilder.newBuilder("testCol1", Types.MinorType.INT.getType()).build());
    schemaBuilder.addField(FieldBuilder.newBuilder("testCol2", Types.MinorType.DATEDAY.getType()).build());
    schemaBuilder.addField(FieldBuilder.newBuilder("testCol3", Types.MinorType.DATEMILLI.getType()).build());
    schemaBuilder.addField(FieldBuilder.newBuilder("testCol4", Types.MinorType.VARBINARY.getType()).build());
    schemaBuilder.addField(FieldBuilder.newBuilder("partition", Types.MinorType.VARCHAR.getType()).build());
    Schema schema = schemaBuilder.build();
    Split split = Mockito.mock(Split.class);
    Mockito.when(split.getProperties()).thenReturn(Collections.singletonMap("partition", "p0"));
    Mockito.when(split.getProperty(Mockito.eq("partition"))).thenReturn("p0");
    Range range1a = Mockito.mock(Range.class, Mockito.RETURNS_DEEP_STUBS);
    Mockito.when(range1a.isSingleValue()).thenReturn(true);
    Mockito.when(range1a.getLow().getValue()).thenReturn(1);
    Range range1b = Mockito.mock(Range.class, Mockito.RETURNS_DEEP_STUBS);
    Mockito.when(range1b.isSingleValue()).thenReturn(true);
    Mockito.when(range1b.getLow().getValue()).thenReturn(2);
    ValueSet valueSet1 = Mockito.mock(SortedRangeSet.class, Mockito.RETURNS_DEEP_STUBS);
    Mockito.when(valueSet1.getRanges().getOrderedRanges()).thenReturn(ImmutableList.of(range1a, range1b));
    final long dateDays = TimeUnit.DAYS.toDays(Date.valueOf("2020-01-05").getTime());
    ValueSet valueSet2 = getSingleValueSet(dateDays);
    Constraints constraints = Mockito.mock(Constraints.class);
    Mockito.when(constraints.getSummary()).thenReturn(new ImmutableMap.Builder<String, ValueSet>().put("testCol2", valueSet2).build());
    PreparedStatement expectedPreparedStatement = Mockito.mock(PreparedStatement.class);
    Mockito.when(this.connection.prepareStatement(Mockito.anyString())).thenReturn(expectedPreparedStatement);
    PreparedStatement preparedStatement = this.hiveRecordHandler.buildSplitSql(this.connection, "testCatalogName", tableName, schema, constraints, split);
    Assert.assertEquals(expectedPreparedStatement, preparedStatement);
}

Example 65 with SchemaBuilder

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

the class ImpalaRecordHandlerTest method buildSplitSql.

@Test
public void buildSplitSql() throws SQLException {
    TableName tableName = new TableName("testSchema", "testTable");
    SchemaBuilder schemaBuilder = SchemaBuilder.newBuilder();
    schemaBuilder.addField(FieldBuilder.newBuilder("testCol1", Types.MinorType.INT.getType()).build());
    schemaBuilder.addField(FieldBuilder.newBuilder("testCol2", Types.MinorType.DATEDAY.getType()).build());
    schemaBuilder.addField(FieldBuilder.newBuilder("testCol3", Types.MinorType.DATEMILLI.getType()).build());
    schemaBuilder.addField(FieldBuilder.newBuilder("testCol4", Types.MinorType.VARBINARY.getType()).build());
    schemaBuilder.addField(FieldBuilder.newBuilder("partition", Types.MinorType.VARCHAR.getType()).build());
    Schema schema = schemaBuilder.build();
    Split split = Mockito.mock(Split.class);
    Mockito.when(split.getProperties()).thenReturn(Collections.singletonMap("partition", "p0"));
    Mockito.when(split.getProperty(Mockito.eq("partition"))).thenReturn("p0");
    Range range1a = Mockito.mock(Range.class, Mockito.RETURNS_DEEP_STUBS);
    Mockito.when(range1a.isSingleValue()).thenReturn(true);
    Mockito.when(range1a.getLow().getValue()).thenReturn(1);
    Range range1b = Mockito.mock(Range.class, Mockito.RETURNS_DEEP_STUBS);
    Mockito.when(range1b.isSingleValue()).thenReturn(true);
    Mockito.when(range1b.getLow().getValue()).thenReturn(2);
    ValueSet valueSet1 = Mockito.mock(SortedRangeSet.class, Mockito.RETURNS_DEEP_STUBS);
    Mockito.when(valueSet1.getRanges().getOrderedRanges()).thenReturn(ImmutableList.of(range1a, range1b));
    final long dateDays = TimeUnit.DAYS.toDays(Date.valueOf("2020-01-05").getTime());
    ValueSet valueSet2 = getSingleValueSet(dateDays);
    Constraints constraints = Mockito.mock(Constraints.class);
    Mockito.when(constraints.getSummary()).thenReturn(new ImmutableMap.Builder<String, ValueSet>().put("testCol2", valueSet2).build());
    PreparedStatement expectedPreparedStatement = Mockito.mock(PreparedStatement.class);
    Mockito.when(this.connection.prepareStatement(Mockito.anyString())).thenReturn(expectedPreparedStatement);
    PreparedStatement preparedStatement = this.impalaRecordHandler.buildSplitSql(this.connection, "testCatalogName", tableName, schema, constraints, split);
    Assert.assertEquals(expectedPreparedStatement, preparedStatement);
}