Examples with Constraints com.amazonaws.athena.connector.lambda.domain.predicate.Constraints used on opensource projects

Example 81 with Constraints

use of com.amazonaws.athena.connector.lambda.domain.predicate.Constraints 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)));
}

Example 82 with Constraints

use of com.amazonaws.athena.connector.lambda.domain.predicate.Constraints in project aws-athena-query-federation by awslabs.

the class DynamoDBMetadataHandlerTest method doGetSplitsScan.

@Test
public void doGetSplitsScan() throws Exception {
    GetTableLayoutResponse layoutResponse = handler.doGetTableLayout(allocator, new GetTableLayoutRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, TEST_TABLE_NAME, new Constraints(ImmutableMap.of()), SchemaBuilder.newBuilder().build(), Collections.EMPTY_SET));
    GetSplitsRequest req = new GetSplitsRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, TEST_TABLE_NAME, layoutResponse.getPartitions(), ImmutableList.of(), new Constraints(new HashMap<>()), null);
    logger.info("doGetSplits: req[{}]", req);
    MetadataResponse rawResponse = handler.doGetSplits(allocator, req);
    assertThat(rawResponse.getRequestType(), equalTo(MetadataRequestType.GET_SPLITS));
    GetSplitsResponse response = (GetSplitsResponse) rawResponse;
    String continuationToken = response.getContinuationToken();
    logger.info("doGetSplits: continuationToken[{}] - numSplits[{}]", continuationToken, response.getSplits().size());
    assertThat(continuationToken == null, is(true));
    Split split = Iterables.getOnlyElement(response.getSplits());
    assertThat(split.getProperty(SEGMENT_ID_PROPERTY), equalTo("0"));
    logger.info("doGetSplitsScan: exit");
}

Example 83 with Constraints

use of com.amazonaws.athena.connector.lambda.domain.predicate.Constraints in project aws-athena-query-federation by awslabs.

the class DynamoDBRecordHandlerTest method testReadScanSplit.

@Test
public void testReadScanSplit() throws Exception {
    Split split = Split.newBuilder(SPILL_LOCATION, keyFactory.create()).add(TABLE_METADATA, TEST_TABLE).add(SEGMENT_ID_PROPERTY, "0").add(SEGMENT_COUNT_METADATA, "1").build();
    ReadRecordsRequest request = new ReadRecordsRequest(TEST_IDENTITY, TEST_CATALOG_NAME, TEST_QUERY_ID, TEST_TABLE_NAME, schema, split, new Constraints(ImmutableMap.of()), // too big to spill
    100_000_000_000L, 100_000_000_000L);
    RecordResponse rawResponse = handler.doReadRecords(allocator, request);
    assertTrue(rawResponse instanceof ReadRecordsResponse);
    ReadRecordsResponse response = (ReadRecordsResponse) rawResponse;
    logger.info("testReadScanSplit: rows[{}]", response.getRecordCount());
    assertEquals(1000, response.getRecords().getRowCount());
    logger.info("testReadScanSplit: {}", BlockUtils.rowToString(response.getRecords(), 0));
}

Example 84 with Constraints

use of com.amazonaws.athena.connector.lambda.domain.predicate.Constraints in project aws-athena-query-federation by awslabs.

the class ElasticsearchRecordHandlerTest method doReadRecordsNoSpill.

@Test
public void doReadRecordsNoSpill() throws Exception {
    logger.info("doReadRecordsNoSpill: enter");
    SearchHit[] searchHit = new SearchHit[2];
    searchHit[0] = new SearchHit(1);
    searchHit[1] = new SearchHit(2);
    SearchHits searchHits = new SearchHits(searchHit, new TotalHits(2, TotalHits.Relation.EQUAL_TO), 4);
    when(mockResponse.getHits()).thenReturn(searchHits);
    Map<String, ValueSet> constraintsMap = new HashMap<>();
    constraintsMap.put("myshort", SortedRangeSet.copyOf(Types.MinorType.SMALLINT.getType(), ImmutableList.of(Range.range(allocator, Types.MinorType.SMALLINT.getType(), (short) 1955, false, (short) 1972, true)), false));
    List<String> expectedProjection = new ArrayList<>();
    mapping.getFields().forEach(field -> expectedProjection.add(field.getName()));
    String expectedPredicate = "(_exists_:myshort) AND myshort:({1955 TO 1972])";
    ReadRecordsRequest request = new ReadRecordsRequest(fakeIdentity(), "elasticsearch", "queryId-" + System.currentTimeMillis(), new TableName("movies", "mishmash"), mapping, split, new Constraints(constraintsMap), // 100GB don't expect this to spill
    100_000_000_000L, 100_000_000_000L);
    RecordResponse rawResponse = handler.doReadRecords(allocator, request);
    // Capture the SearchRequest object from the call to client.getDocuments().
    // The former contains information such as the projection and predicate.
    ArgumentCaptor<SearchRequest> argumentCaptor = ArgumentCaptor.forClass(SearchRequest.class);
    verify(mockClient).getDocuments(argumentCaptor.capture());
    SearchRequest searchRequest = argumentCaptor.getValue();
    // Get the actual projection and compare to the expected one.
    List<String> actualProjection = ImmutableList.copyOf(searchRequest.source().fetchSource().includes());
    assertEquals("Projections do not match", expectedProjection, actualProjection);
    // Get the actual predicate and compare to the expected one.
    String actualPredicate = searchRequest.source().query().queryName();
    assertEquals("Predicates do not match", expectedPredicate, actualPredicate);
    assertTrue(rawResponse instanceof ReadRecordsResponse);
    ReadRecordsResponse response = (ReadRecordsResponse) rawResponse;
    logger.info("doReadRecordsNoSpill: rows[{}]", response.getRecordCount());
    assertEquals(2, response.getRecords().getRowCount());
    for (int i = 0; i < response.getRecords().getRowCount(); ++i) {
        logger.info("doReadRecordsNoSpill - Row: {}, {}", i, BlockUtils.rowToString(response.getRecords(), i));
    }
    logger.info("doReadRecordsNoSpill: exit");
}

Example 85 with Constraints

use of com.amazonaws.athena.connector.lambda.domain.predicate.Constraints in project aws-athena-query-federation by awslabs.

the class DynamoDBRecordHandlerTest method testReadQuerySplit.

@Test
public void testReadQuerySplit() throws Exception {
    Map<String, String> expressionNames = ImmutableMap.of("#col_1", "col_1");
    Map<String, AttributeValue> expressionValues = ImmutableMap.of(":v0", toAttributeValue(1));
    Split split = Split.newBuilder(SPILL_LOCATION, keyFactory.create()).add(TABLE_METADATA, TEST_TABLE).add(HASH_KEY_NAME_METADATA, "col_0").add("col_0", toJsonString(toAttributeValue("test_str_0"))).add(RANGE_KEY_FILTER_METADATA, "#col_1 >= :v0").add(EXPRESSION_NAMES_METADATA, toJsonString(expressionNames)).add(EXPRESSION_VALUES_METADATA, toJsonString(expressionValues)).build();
    ReadRecordsRequest request = new ReadRecordsRequest(TEST_IDENTITY, TEST_CATALOG_NAME, TEST_QUERY_ID, TEST_TABLE_NAME, schema, split, new Constraints(ImmutableMap.of()), // too big to spill
    100_000_000_000L, 100_000_000_000L);
    RecordResponse rawResponse = handler.doReadRecords(allocator, request);
    assertTrue(rawResponse instanceof ReadRecordsResponse);
    ReadRecordsResponse response = (ReadRecordsResponse) rawResponse;
    logger.info("testReadQuerySplit: rows[{}]", response.getRecordCount());
    assertEquals(2, response.getRecords().getRowCount());
    logger.info("testReadQuerySplit: {}", BlockUtils.rowToString(response.getRecords(), 0));
}