Search in sources :

Example 31 with RecordResponse

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

the class MetricsRecordHandlerTest method readMetricsWithConstraint.

@Test
public void readMetricsWithConstraint() throws Exception {
    logger.info("readMetricsWithConstraint: enter");
    String namespace = "namespace";
    String dimName = "dimName";
    String dimValue = "dimValye";
    int numMetrics = 100;
    AtomicLong numCalls = new AtomicLong(0);
    when(mockMetrics.listMetrics(any(ListMetricsRequest.class))).thenAnswer((InvocationOnMock invocation) -> {
        ListMetricsRequest request = invocation.getArgumentAt(0, ListMetricsRequest.class);
        numCalls.incrementAndGet();
        // assert that the namespace filter was indeed pushed down
        assertEquals(namespace, request.getNamespace());
        String nextToken = (request.getNextToken() == null) ? "valid" : null;
        List<Metric> metrics = new ArrayList<>();
        for (int i = 0; i < numMetrics; i++) {
            metrics.add(new Metric().withNamespace(namespace).withMetricName("metric-" + i).withDimensions(new Dimension().withName(dimName).withValue(dimValue)));
            metrics.add(new Metric().withNamespace(namespace + i).withMetricName("metric-" + i));
        }
        return new ListMetricsResult().withNextToken(nextToken).withMetrics(metrics);
    });
    Map<String, ValueSet> constraintsMap = new HashMap<>();
    constraintsMap.put(NAMESPACE_FIELD, makeStringEquals(allocator, namespace));
    constraintsMap.put(DIMENSION_NAME_FIELD, makeStringEquals(allocator, dimName));
    constraintsMap.put(DIMENSION_VALUE_FIELD, makeStringEquals(allocator, dimValue));
    S3SpillLocation spillLocation = S3SpillLocation.newBuilder().withBucket(UUID.randomUUID().toString()).withSplitId(UUID.randomUUID().toString()).withQueryId(UUID.randomUUID().toString()).withIsDirectory(true).build();
    Split split = Split.newBuilder(spillLocation, keyFactory.create()).build();
    ReadRecordsRequest request = new ReadRecordsRequest(identity, "catalog", "queryId-" + System.currentTimeMillis(), METRICS_TABLE_NAME, METRIC_TABLE.getSchema(), split, new Constraints(constraintsMap), 100_000_000_000L, // 100GB don't expect this to spill
    100_000_000_000L);
    RecordResponse rawResponse = handler.doReadRecords(allocator, request);
    assertTrue(rawResponse instanceof ReadRecordsResponse);
    ReadRecordsResponse response = (ReadRecordsResponse) rawResponse;
    logger.info("readMetricsWithConstraint: rows[{}]", response.getRecordCount());
    assertEquals(numCalls.get() * numMetrics, response.getRecords().getRowCount());
    logger.info("readMetricsWithConstraint: {}", BlockUtils.rowToString(response.getRecords(), 0));
    logger.info("readMetricsWithConstraint: exit");
}
Also used : HashMap(java.util.HashMap) ReadRecordsResponse(com.amazonaws.athena.connector.lambda.records.ReadRecordsResponse) ArrayList(java.util.ArrayList) ListMetricsResult(com.amazonaws.services.cloudwatch.model.ListMetricsResult) Matchers.anyString(org.mockito.Matchers.anyString) Dimension(com.amazonaws.services.cloudwatch.model.Dimension) RecordResponse(com.amazonaws.athena.connector.lambda.records.RecordResponse) AtomicLong(java.util.concurrent.atomic.AtomicLong) ReadRecordsRequest(com.amazonaws.athena.connector.lambda.records.ReadRecordsRequest) Constraints(com.amazonaws.athena.connector.lambda.domain.predicate.Constraints) InvocationOnMock(org.mockito.invocation.InvocationOnMock) S3SpillLocation(com.amazonaws.athena.connector.lambda.domain.spill.S3SpillLocation) ListMetricsRequest(com.amazonaws.services.cloudwatch.model.ListMetricsRequest) Metric(com.amazonaws.services.cloudwatch.model.Metric) Split(com.amazonaws.athena.connector.lambda.domain.Split) ValueSet(com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet) Test(org.junit.Test)

Example 32 with RecordResponse

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

the class CloudwatchRecordHandlerTest method doReadRecordsNoSpill.

@Test
public void doReadRecordsNoSpill() throws Exception {
    logger.info("doReadRecordsNoSpill: enter");
    Map<String, ValueSet> constraintsMap = new HashMap<>();
    constraintsMap.put("time", SortedRangeSet.copyOf(Types.MinorType.BIGINT.getType(), ImmutableList.of(Range.equal(allocator, Types.MinorType.BIGINT.getType(), 100L)), false));
    ReadRecordsRequest request = new ReadRecordsRequest(identity, "catalog", "queryId-" + System.currentTimeMillis(), new TableName("schema", "table"), schemaForRead, Split.newBuilder(S3SpillLocation.newBuilder().withBucket(UUID.randomUUID().toString()).withSplitId(UUID.randomUUID().toString()).withQueryId(UUID.randomUUID().toString()).withIsDirectory(true).build(), keyFactory.create()).add(CloudwatchMetadataHandler.LOG_STREAM_FIELD, "table").build(), new Constraints(constraintsMap), 100_000_000_000L, // 100GB don't expect this to spill
    100_000_000_000L);
    RecordResponse rawResponse = handler.doReadRecords(allocator, request);
    assertTrue(rawResponse instanceof ReadRecordsResponse);
    ReadRecordsResponse response = (ReadRecordsResponse) rawResponse;
    logger.info("doReadRecordsNoSpill: rows[{}]", response.getRecordCount());
    assertTrue(response.getRecords().getRowCount() == 3);
    logger.info("doReadRecordsNoSpill: {}", BlockUtils.rowToString(response.getRecords(), 0));
    logger.info("doReadRecordsNoSpill: exit");
}
Also used : TableName(com.amazonaws.athena.connector.lambda.domain.TableName) ReadRecordsRequest(com.amazonaws.athena.connector.lambda.records.ReadRecordsRequest) Constraints(com.amazonaws.athena.connector.lambda.domain.predicate.Constraints) HashMap(java.util.HashMap) ReadRecordsResponse(com.amazonaws.athena.connector.lambda.records.ReadRecordsResponse) RemoteReadRecordsResponse(com.amazonaws.athena.connector.lambda.records.RemoteReadRecordsResponse) Matchers.anyString(org.mockito.Matchers.anyString) RecordResponse(com.amazonaws.athena.connector.lambda.records.RecordResponse) ValueSet(com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet) Test(org.junit.Test)

Example 33 with RecordResponse

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

the class DynamoDBRecordHandlerTest method testStructWithNullFromDdbTable.

@Test
public void testStructWithNullFromDdbTable() throws Exception {
    when(glueClient.getTable(any())).thenThrow(new EntityNotFoundException(""));
    TableName tableName = new TableName(DEFAULT_SCHEMA, TEST_TABLE4);
    GetTableRequest getTableRequest = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, tableName);
    GetTableResponse getTableResponse = metadataHandler.doGetTable(allocator, getTableRequest);
    logger.info("testStructWithNullFromGlueTable: GetTableResponse[{}]", getTableResponse);
    logger.info("testStructWithNullFromGlueTable: GetTableResponse Schema[{}]", getTableResponse.getSchema());
    Schema schema4 = getTableResponse.getSchema();
    for (Field f : schema4.getFields()) {
        if (f.getName().equals("Col2")) {
            assertEquals(1, f.getChildren().size());
            assertTrue(f.getType() instanceof ArrowType.Struct);
        }
    }
    Split split = Split.newBuilder(SPILL_LOCATION, keyFactory.create()).add(TABLE_METADATA, TEST_TABLE4).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_4_NAME, schema4, 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("testStructWithNullFromGlueTable: {}", BlockUtils.rowToString(response.getRecords(), 0));
    Block result = response.getRecords();
    assertEquals(1, result.getRowCount());
    assertEquals(schema4, result.getSchema());
    assertEquals("[Col0 : hashVal], [Col1 : {[field1 : someField1]}]", BlockUtils.rowToString(response.getRecords(), 0));
}
Also used : ReadRecordsResponse(com.amazonaws.athena.connector.lambda.records.ReadRecordsResponse) Schema(org.apache.arrow.vector.types.pojo.Schema) ArrowType(org.apache.arrow.vector.types.pojo.ArrowType) EntityNotFoundException(com.amazonaws.services.glue.model.EntityNotFoundException) RecordResponse(com.amazonaws.athena.connector.lambda.records.RecordResponse) TableName(com.amazonaws.athena.connector.lambda.domain.TableName) GetTableRequest(com.amazonaws.athena.connector.lambda.metadata.GetTableRequest) Field(org.apache.arrow.vector.types.pojo.Field) ReadRecordsRequest(com.amazonaws.athena.connector.lambda.records.ReadRecordsRequest) Constraints(com.amazonaws.athena.connector.lambda.domain.predicate.Constraints) GetTableResponse(com.amazonaws.athena.connector.lambda.metadata.GetTableResponse) Block(com.amazonaws.athena.connector.lambda.data.Block) Split(com.amazonaws.athena.connector.lambda.domain.Split) Test(org.junit.Test)

Example 34 with RecordResponse

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

the class DynamoDBRecordHandlerTest method testZeroRowQuery.

@Test
public void testZeroRowQuery() 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_999999"))).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("testZeroRowQuery: rows[{}]", response.getRecordCount());
    assertEquals(0, response.getRecords().getRowCount());
}
Also used : ItemUtils.toAttributeValue(com.amazonaws.services.dynamodbv2.document.ItemUtils.toAttributeValue) AttributeValue(com.amazonaws.services.dynamodbv2.model.AttributeValue) ReadRecordsRequest(com.amazonaws.athena.connector.lambda.records.ReadRecordsRequest) Constraints(com.amazonaws.athena.connector.lambda.domain.predicate.Constraints) ReadRecordsResponse(com.amazonaws.athena.connector.lambda.records.ReadRecordsResponse) Jackson.toJsonString(com.amazonaws.util.json.Jackson.toJsonString) RecordResponse(com.amazonaws.athena.connector.lambda.records.RecordResponse) Split(com.amazonaws.athena.connector.lambda.domain.Split) Test(org.junit.Test)

Example 35 with RecordResponse

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

the class DynamoDBRecordHandlerTest method testReadScanSplitFiltered.

@Test
public void testReadScanSplitFiltered() throws Exception {
    Map<String, String> expressionNames = ImmutableMap.of("#col_6", "col_6");
    Map<String, AttributeValue> expressionValues = ImmutableMap.of(":v0", toAttributeValue(0), ":v1", toAttributeValue(1));
    Split split = Split.newBuilder(SPILL_LOCATION, keyFactory.create()).add(TABLE_METADATA, TEST_TABLE).add(SEGMENT_ID_PROPERTY, "0").add(SEGMENT_COUNT_METADATA, "1").add(NON_KEY_FILTER_METADATA, "NOT #col_6 IN (:v0,:v1)").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("testReadScanSplitFiltered: rows[{}]", response.getRecordCount());
    assertEquals(992, response.getRecords().getRowCount());
    logger.info("testReadScanSplitFiltered: {}", BlockUtils.rowToString(response.getRecords(), 0));
}
Also used : ItemUtils.toAttributeValue(com.amazonaws.services.dynamodbv2.document.ItemUtils.toAttributeValue) AttributeValue(com.amazonaws.services.dynamodbv2.model.AttributeValue) ReadRecordsRequest(com.amazonaws.athena.connector.lambda.records.ReadRecordsRequest) Constraints(com.amazonaws.athena.connector.lambda.domain.predicate.Constraints) ReadRecordsResponse(com.amazonaws.athena.connector.lambda.records.ReadRecordsResponse) Jackson.toJsonString(com.amazonaws.util.json.Jackson.toJsonString) RecordResponse(com.amazonaws.athena.connector.lambda.records.RecordResponse) Split(com.amazonaws.athena.connector.lambda.domain.Split) Test(org.junit.Test)

Aggregations

RecordResponse (com.amazonaws.athena.connector.lambda.records.RecordResponse)37 Constraints (com.amazonaws.athena.connector.lambda.domain.predicate.Constraints)33 ReadRecordsRequest (com.amazonaws.athena.connector.lambda.records.ReadRecordsRequest)33 Test (org.junit.Test)32 ReadRecordsResponse (com.amazonaws.athena.connector.lambda.records.ReadRecordsResponse)27 ValueSet (com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet)24 HashMap (java.util.HashMap)24 Matchers.anyString (org.mockito.Matchers.anyString)22 RemoteReadRecordsResponse (com.amazonaws.athena.connector.lambda.records.RemoteReadRecordsResponse)21 S3SpillLocation (com.amazonaws.athena.connector.lambda.domain.spill.S3SpillLocation)20 TableName (com.amazonaws.athena.connector.lambda.domain.TableName)19 Split (com.amazonaws.athena.connector.lambda.domain.Split)18 ArrayList (java.util.ArrayList)18 InvocationOnMock (org.mockito.invocation.InvocationOnMock)14 Block (com.amazonaws.athena.connector.lambda.data.Block)13 SpillLocation (com.amazonaws.athena.connector.lambda.domain.spill.SpillLocation)11 Schema (org.apache.arrow.vector.types.pojo.Schema)10 EquatableValueSet (com.amazonaws.athena.connector.lambda.domain.predicate.EquatableValueSet)9 InputStream (java.io.InputStream)7 List (java.util.List)7