Search in sources :

Example 1 with ItemUtils.toAttributeValue

use of com.amazonaws.services.dynamodbv2.document.ItemUtils.toAttributeValue in project aws-athena-query-federation by awslabs.

the class DynamoDBMetadataHandlerTest method doGetTableLayoutQueryIndex.

@Test
public void doGetTableLayoutQueryIndex() throws Exception {
    Map<String, ValueSet> constraintsMap = new HashMap<>();
    SortedRangeSet.Builder dateValueSet = SortedRangeSet.newBuilder(Types.MinorType.DATEDAY.getType(), false);
    SortedRangeSet.Builder timeValueSet = SortedRangeSet.newBuilder(Types.MinorType.DATEMILLI.getType(), false);
    LocalDateTime dateTime = LocalDateTime.of(2019, 9, 23, 11, 18, 37);
    // Set to Epoch time
    Instant epoch = Instant.MIN;
    dateValueSet.add(Range.equal(allocator, Types.MinorType.DATEDAY.getType(), ChronoUnit.DAYS.between(epoch, dateTime.toInstant(ZoneOffset.UTC))));
    LocalDateTime dateTime2 = dateTime.plusHours(26);
    dateValueSet.add(Range.equal(allocator, Types.MinorType.DATEDAY.getType(), ChronoUnit.DAYS.between(epoch, dateTime2.toInstant(ZoneOffset.UTC))));
    long startTime = dateTime.toInstant(ZoneOffset.UTC).toEpochMilli();
    long endTime = dateTime2.toInstant(ZoneOffset.UTC).toEpochMilli();
    timeValueSet.add(Range.range(allocator, Types.MinorType.DATEMILLI.getType(), startTime, true, endTime, true));
    constraintsMap.put("col_4", dateValueSet.build());
    constraintsMap.put("col_5", timeValueSet.build());
    GetTableLayoutResponse res = handler.doGetTableLayout(allocator, new GetTableLayoutRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, TEST_TABLE_NAME, new Constraints(constraintsMap), SchemaBuilder.newBuilder().build(), Collections.EMPTY_SET));
    logger.info("doGetTableLayout schema - {}", res.getPartitions().getSchema());
    logger.info("doGetTableLayout partitions - {}", res.getPartitions());
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(PARTITION_TYPE_METADATA), equalTo(QUERY_PARTITION_TYPE));
    assertThat(res.getPartitions().getSchema().getCustomMetadata().containsKey(INDEX_METADATA), is(true));
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(INDEX_METADATA), equalTo("test_index"));
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(HASH_KEY_NAME_METADATA), equalTo("col_4"));
    assertThat(res.getPartitions().getRowCount(), equalTo(2));
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(RANGE_KEY_NAME_METADATA), equalTo("col_5"));
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(RANGE_KEY_FILTER_METADATA), equalTo("(#col_5 >= :v0 AND #col_5 <= :v1)"));
    ImmutableMap<String, String> expressionNames = ImmutableMap.of("#col_4", "col_4", "#col_5", "col_5");
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(EXPRESSION_NAMES_METADATA), equalTo(Jackson.toJsonString(expressionNames)));
    ImmutableMap<String, AttributeValue> expressionValues = ImmutableMap.of(":v0", ItemUtils.toAttributeValue(startTime), ":v1", ItemUtils.toAttributeValue(endTime));
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(EXPRESSION_VALUES_METADATA), equalTo(Jackson.toJsonString(expressionValues)));
}
Also used : LocalDateTime(java.time.LocalDateTime) AttributeValue(com.amazonaws.services.dynamodbv2.model.AttributeValue) HashMap(java.util.HashMap) Instant(java.time.Instant) Constraints(com.amazonaws.athena.connector.lambda.domain.predicate.Constraints) SortedRangeSet(com.amazonaws.athena.connector.lambda.domain.predicate.SortedRangeSet) GetTableLayoutResponse(com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutResponse) GetTableLayoutRequest(com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest) ValueSet(com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet) EquatableValueSet(com.amazonaws.athena.connector.lambda.domain.predicate.EquatableValueSet) Test(org.junit.Test)

Example 2 with ItemUtils.toAttributeValue

use of com.amazonaws.services.dynamodbv2.document.ItemUtils.toAttributeValue in project aws-athena-query-federation by awslabs.

the class DynamoDBMetadataHandlerTest method doGetTableLayoutScanWithTypeOverride.

@Test
public void doGetTableLayoutScanWithTypeOverride() throws Exception {
    List<Column> columns = new ArrayList<>();
    columns.add(new Column().withName("col1").withType("int"));
    columns.add(new Column().withName("col2").withType("timestamptz"));
    columns.add(new Column().withName("col3").withType("string"));
    Map<String, String> param = ImmutableMap.of(SOURCE_TABLE_PROPERTY, TEST_TABLE, COLUMN_NAME_MAPPING_PROPERTY, "col1=Col1", DATETIME_FORMAT_MAPPING_PROPERTY, "col1=datetime1,col3=datetime3 ");
    Table table = new Table().withParameters(param).withPartitionKeys().withStorageDescriptor(new StorageDescriptor().withColumns(columns));
    GetTableResult mockResult = new GetTableResult().withTable(table);
    when(glueClient.getTable(any())).thenReturn(mockResult);
    TableName tableName = new TableName(DEFAULT_SCHEMA, "glueTableForTestTable");
    GetTableRequest getTableRequest = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, tableName);
    GetTableResponse getTableResponse = handler.doGetTable(allocator, getTableRequest);
    logger.info("validateSourceTableNamePropagation: GetTableResponse[{}]", getTableResponse);
    Map<String, String> customMetadata = getTableResponse.getSchema().getCustomMetadata();
    assertThat(customMetadata.get(SOURCE_TABLE_PROPERTY), equalTo(TEST_TABLE));
    assertThat(customMetadata.get(DATETIME_FORMAT_MAPPING_PROPERTY_NORMALIZED), equalTo("Col1=datetime1,col3=datetime3"));
    Map<String, ValueSet> constraintsMap = new HashMap<>();
    constraintsMap.put("col3", EquatableValueSet.newBuilder(allocator, new ArrowType.Bool(), true, true).add(true).build());
    constraintsMap.put("col2", EquatableValueSet.newBuilder(allocator, new ArrowType.Bool(), true, true).add(true).build());
    GetTableLayoutRequest getTableLayoutRequest = new GetTableLayoutRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, tableName, new Constraints(constraintsMap), getTableResponse.getSchema(), Collections.EMPTY_SET);
    GetTableLayoutResponse res = handler.doGetTableLayout(allocator, getTableLayoutRequest);
    logger.info("doGetTableLayoutScanWithTypeOverride schema - {}", res.getPartitions().getSchema());
    logger.info("doGetTableLayoutScanWithTypeOverride partitions - {}", res.getPartitions());
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(PARTITION_TYPE_METADATA), equalTo(SCAN_PARTITION_TYPE));
    // no hash key constraints, so look for segment count column
    assertThat(res.getPartitions().getSchema().findField(SEGMENT_COUNT_METADATA) != null, is(true));
    assertThat(res.getPartitions().getRowCount(), equalTo(1));
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(NON_KEY_FILTER_METADATA), equalTo("(#col3 = :v0 OR attribute_not_exists(#col3) OR #col3 = :v1)"));
    ImmutableMap<String, String> expressionNames = ImmutableMap.of("#col3", "col3", "#col2", "col2");
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(EXPRESSION_NAMES_METADATA), equalTo(Jackson.toJsonString(expressionNames)));
    ImmutableMap<String, AttributeValue> expressionValues = ImmutableMap.of(":v0", ItemUtils.toAttributeValue(true), ":v1", ItemUtils.toAttributeValue(null));
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(EXPRESSION_VALUES_METADATA), equalTo(Jackson.toJsonString(expressionValues)));
}
Also used : AttributeValue(com.amazonaws.services.dynamodbv2.model.AttributeValue) Table(com.amazonaws.services.glue.model.Table) HashMap(java.util.HashMap) ArrayList(java.util.ArrayList) StorageDescriptor(com.amazonaws.services.glue.model.StorageDescriptor) TableName(com.amazonaws.athena.connector.lambda.domain.TableName) GetTableRequest(com.amazonaws.athena.connector.lambda.metadata.GetTableRequest) Constraints(com.amazonaws.athena.connector.lambda.domain.predicate.Constraints) GetTableLayoutResponse(com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutResponse) Column(com.amazonaws.services.glue.model.Column) GetTableResponse(com.amazonaws.athena.connector.lambda.metadata.GetTableResponse) GetTableLayoutRequest(com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest) ValueSet(com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet) EquatableValueSet(com.amazonaws.athena.connector.lambda.domain.predicate.EquatableValueSet) GetTableResult(com.amazonaws.services.glue.model.GetTableResult) Test(org.junit.Test)

Example 3 with ItemUtils.toAttributeValue

use of com.amazonaws.services.dynamodbv2.document.ItemUtils.toAttributeValue in project aws-athena-query-federation by awslabs.

the class DynamoDBMetadataHandlerTest method doGetTableLayoutScan.

@Test
public void doGetTableLayoutScan() throws Exception {
    Map<String, ValueSet> constraintsMap = new HashMap<>();
    constraintsMap.put("col_3", EquatableValueSet.newBuilder(allocator, new ArrowType.Bool(), true, true).add(true).build());
    GetTableLayoutRequest req = new GetTableLayoutRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, new TableName(TEST_CATALOG_NAME, TEST_TABLE), new Constraints(constraintsMap), SchemaBuilder.newBuilder().build(), Collections.EMPTY_SET);
    GetTableLayoutResponse res = handler.doGetTableLayout(allocator, req);
    logger.info("doGetTableLayout schema - {}", res.getPartitions().getSchema());
    logger.info("doGetTableLayout partitions - {}", res.getPartitions());
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(PARTITION_TYPE_METADATA), equalTo(SCAN_PARTITION_TYPE));
    // no hash key constraints, so look for segment count column
    assertThat(res.getPartitions().getSchema().findField(SEGMENT_COUNT_METADATA) != null, is(true));
    assertThat(res.getPartitions().getRowCount(), equalTo(1));
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(NON_KEY_FILTER_METADATA), equalTo("(#col_3 = :v0 OR attribute_not_exists(#col_3) OR #col_3 = :v1)"));
    ImmutableMap<String, String> expressionNames = ImmutableMap.of("#col_3", "col_3");
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(EXPRESSION_NAMES_METADATA), equalTo(Jackson.toJsonString(expressionNames)));
    ImmutableMap<String, AttributeValue> expressionValues = ImmutableMap.of(":v0", ItemUtils.toAttributeValue(true), ":v1", ItemUtils.toAttributeValue(null));
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(EXPRESSION_VALUES_METADATA), equalTo(Jackson.toJsonString(expressionValues)));
}
Also used : TableName(com.amazonaws.athena.connector.lambda.domain.TableName) AttributeValue(com.amazonaws.services.dynamodbv2.model.AttributeValue) Constraints(com.amazonaws.athena.connector.lambda.domain.predicate.Constraints) GetTableLayoutResponse(com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutResponse) HashMap(java.util.HashMap) GetTableLayoutRequest(com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest) ValueSet(com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet) EquatableValueSet(com.amazonaws.athena.connector.lambda.domain.predicate.EquatableValueSet) Test(org.junit.Test)

Aggregations

Constraints (com.amazonaws.athena.connector.lambda.domain.predicate.Constraints)3 EquatableValueSet (com.amazonaws.athena.connector.lambda.domain.predicate.EquatableValueSet)3 ValueSet (com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet)3 GetTableLayoutRequest (com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest)3 GetTableLayoutResponse (com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutResponse)3 AttributeValue (com.amazonaws.services.dynamodbv2.model.AttributeValue)3 HashMap (java.util.HashMap)3 Test (org.junit.Test)3 TableName (com.amazonaws.athena.connector.lambda.domain.TableName)2 SortedRangeSet (com.amazonaws.athena.connector.lambda.domain.predicate.SortedRangeSet)1 GetTableRequest (com.amazonaws.athena.connector.lambda.metadata.GetTableRequest)1 GetTableResponse (com.amazonaws.athena.connector.lambda.metadata.GetTableResponse)1 Column (com.amazonaws.services.glue.model.Column)1 GetTableResult (com.amazonaws.services.glue.model.GetTableResult)1 StorageDescriptor (com.amazonaws.services.glue.model.StorageDescriptor)1 Table (com.amazonaws.services.glue.model.Table)1 Instant (java.time.Instant)1 LocalDateTime (java.time.LocalDateTime)1 ArrayList (java.util.ArrayList)1