use of com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet in project aws-athena-query-federation by awslabs.
the class MetricsRecordHandler method readMetricSamplesWithConstraint.
/**
* Handles retrieving the samples for a specific metric from Cloudwatch Metrics.
*/
private void readMetricSamplesWithConstraint(BlockSpiller blockSpiller, ReadRecordsRequest request, QueryStatusChecker queryStatusChecker) throws TimeoutException {
GetMetricDataRequest dataRequest = MetricUtils.makeGetMetricDataRequest(request);
Map<String, MetricDataQuery> queries = new HashMap<>();
for (MetricDataQuery query : dataRequest.getMetricDataQueries()) {
queries.put(query.getId(), query);
}
String prevToken;
ValueSet dimensionNameConstraint = request.getConstraints().getSummary().get(DIMENSION_NAME_FIELD);
ValueSet dimensionValueConstraint = request.getConstraints().getSummary().get(DIMENSION_VALUE_FIELD);
do {
prevToken = dataRequest.getNextToken();
GetMetricDataResult result = invoker.invoke(() -> metrics.getMetricData(dataRequest));
for (MetricDataResult nextMetric : result.getMetricDataResults()) {
MetricStat metricStat = queries.get(nextMetric.getId()).getMetricStat();
List<Date> timestamps = nextMetric.getTimestamps();
List<Double> values = nextMetric.getValues();
for (int i = 0; i < nextMetric.getValues().size(); i++) {
int sampleNum = i;
blockSpiller.writeRows((Block block, int row) -> {
/**
* Most constraints were already applied at split generation so we only need to apply
* a subset.
*/
block.offerValue(METRIC_NAME_FIELD, row, metricStat.getMetric().getMetricName());
block.offerValue(NAMESPACE_FIELD, row, metricStat.getMetric().getNamespace());
block.offerValue(STATISTIC_FIELD, row, metricStat.getStat());
block.offerComplexValue(DIMENSIONS_FIELD, row, (Field field, Object val) -> {
if (field.getName().equals(DIMENSION_NAME_FIELD)) {
return ((Dimension) val).getName();
} else if (field.getName().equals(DIMENSION_VALUE_FIELD)) {
return ((Dimension) val).getValue();
}
throw new RuntimeException("Unexpected field " + field.getName());
}, metricStat.getMetric().getDimensions());
// This field is 'faked' in that we just use it as a convenient way to filter single dimensions. As such
// we always populate it with the value of the filter if the constraint passed and the filter was singleValue
String dimName = (dimensionNameConstraint == null || !dimensionNameConstraint.isSingleValue()) ? null : dimensionNameConstraint.getSingleValue().toString();
block.offerValue(DIMENSION_NAME_FIELD, row, dimName);
// This field is 'faked' in that we just use it as a convenient way to filter single dimensions. As such
// we always populate it with the value of the filter if the constraint passed and the filter was singleValue
String dimVal = (dimensionValueConstraint == null || !dimensionValueConstraint.isSingleValue()) ? null : dimensionValueConstraint.getSingleValue().toString();
block.offerValue(DIMENSION_VALUE_FIELD, row, dimVal);
block.offerValue(PERIOD_FIELD, row, metricStat.getPeriod());
boolean matches = true;
block.offerValue(VALUE_FIELD, row, values.get(sampleNum));
long timestamp = timestamps.get(sampleNum).getTime() / 1000;
block.offerValue(TIMESTAMP_FIELD, row, timestamp);
return matches ? 1 : 0;
});
}
}
dataRequest.setNextToken(result.getNextToken());
} while (dataRequest.getNextToken() != null && !dataRequest.getNextToken().equalsIgnoreCase(prevToken) && queryStatusChecker.isQueryRunning());
}
use of com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet in project aws-athena-query-federation by awslabs.
the class MetricsRecordHandler method readMetricsWithConstraint.
/**
* Handles retrieving the list of available metrics when the METRICS_TABLE is queried by listing metrics in Cloudwatch Metrics.
*/
private void readMetricsWithConstraint(BlockSpiller blockSpiller, ReadRecordsRequest request, QueryStatusChecker queryStatusChecker) throws TimeoutException {
ListMetricsRequest listMetricsRequest = new ListMetricsRequest();
MetricUtils.pushDownPredicate(request.getConstraints(), listMetricsRequest);
String prevToken;
Set<String> requiredFields = new HashSet<>();
request.getSchema().getFields().stream().forEach(next -> requiredFields.add(next.getName()));
ValueSet dimensionNameConstraint = request.getConstraints().getSummary().get(DIMENSION_NAME_FIELD);
ValueSet dimensionValueConstraint = request.getConstraints().getSummary().get(DIMENSION_VALUE_FIELD);
do {
prevToken = listMetricsRequest.getNextToken();
ListMetricsResult result = invoker.invoke(() -> metrics.listMetrics(listMetricsRequest));
for (Metric nextMetric : result.getMetrics()) {
blockSpiller.writeRows((Block block, int row) -> {
boolean matches = MetricUtils.applyMetricConstraints(blockSpiller.getConstraintEvaluator(), nextMetric, null);
if (matches) {
matches &= block.offerValue(METRIC_NAME_FIELD, row, nextMetric.getMetricName());
matches &= block.offerValue(NAMESPACE_FIELD, row, nextMetric.getNamespace());
matches &= block.offerComplexValue(STATISTIC_FIELD, row, DEFAULT, STATISTICS);
matches &= block.offerComplexValue(DIMENSIONS_FIELD, row, (Field field, Object val) -> {
if (field.getName().equals(DIMENSION_NAME_FIELD)) {
return ((Dimension) val).getName();
} else if (field.getName().equals(DIMENSION_VALUE_FIELD)) {
return ((Dimension) val).getValue();
}
throw new RuntimeException("Unexpected field " + field.getName());
}, nextMetric.getDimensions());
// This field is 'faked' in that we just use it as a convenient way to filter single dimensions. As such
// we always populate it with the value of the filter if the constraint passed and the filter was singleValue
String dimName = (dimensionNameConstraint == null || !dimensionNameConstraint.isSingleValue()) ? null : (dimensionNameConstraint.getSingleValue().toString());
matches &= block.offerValue(DIMENSION_NAME_FIELD, row, dimName);
// This field is 'faked' in that we just use it as a convenient way to filter single dimensions. As such
// we always populate it with the value of the filter if the constraint passed and the filter was singleValue
String dimValue = (dimensionValueConstraint == null || !dimensionValueConstraint.isSingleValue()) ? null : dimensionValueConstraint.getSingleValue().toString();
matches &= block.offerValue(DIMENSION_VALUE_FIELD, row, dimValue);
}
return matches ? 1 : 0;
});
}
listMetricsRequest.setNextToken(result.getNextToken());
} while (listMetricsRequest.getNextToken() != null && !listMetricsRequest.getNextToken().equalsIgnoreCase(prevToken) && queryStatusChecker.isQueryRunning());
}
use of com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet in project aws-athena-query-federation by awslabs.
the class MetricUtilsTest method applyMetricConstraints.
@Test
public void applyMetricConstraints() {
Schema schema = SchemaBuilder.newBuilder().addStringField(NAMESPACE_FIELD).addStringField(METRIC_NAME_FIELD).addStringField(STATISTIC_FIELD).addStringField(DIMENSION_NAME_FIELD).addStringField(DIMENSION_VALUE_FIELD).build();
Map<String, ValueSet> constraintsMap = new HashMap<>();
constraintsMap.put(NAMESPACE_FIELD, makeStringEquals(allocator, "match1"));
constraintsMap.put(METRIC_NAME_FIELD, makeStringEquals(allocator, "match2"));
constraintsMap.put(STATISTIC_FIELD, makeStringEquals(allocator, "match3"));
constraintsMap.put(DIMENSION_NAME_FIELD, makeStringEquals(allocator, "match4"));
constraintsMap.put(DIMENSION_VALUE_FIELD, makeStringEquals(allocator, "match5"));
ConstraintEvaluator constraintEvaluator = new ConstraintEvaluator(allocator, schema, new Constraints(constraintsMap));
Metric metric = new Metric().withNamespace("match1").withMetricName("match2").withDimensions(new Dimension().withName("match4").withValue("match5"));
String statistic = "match3";
assertTrue(MetricUtils.applyMetricConstraints(constraintEvaluator, metric, statistic));
assertFalse(MetricUtils.applyMetricConstraints(constraintEvaluator, copyMetric(metric).withNamespace("no_match"), statistic));
assertFalse(MetricUtils.applyMetricConstraints(constraintEvaluator, copyMetric(metric).withMetricName("no_match"), statistic));
assertFalse(MetricUtils.applyMetricConstraints(constraintEvaluator, copyMetric(metric).withDimensions(Collections.singletonList(new Dimension().withName("no_match").withValue("match5"))), statistic));
assertFalse(MetricUtils.applyMetricConstraints(constraintEvaluator, copyMetric(metric).withDimensions(Collections.singletonList(new Dimension().withName("match4").withValue("no_match"))), statistic));
assertFalse(MetricUtils.applyMetricConstraints(constraintEvaluator, copyMetric(metric), "no_match"));
}
use of com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet in project aws-athena-query-federation by awslabs.
the class DataLakeRecordHandlerTest method buildSplitSqlNew.
@Test
public void buildSplitSqlNew() 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.VARCHAR.getType()).build());
Schema schema = schemaBuilder.build();
Split split = Mockito.mock(Split.class);
Mockito.when(split.getProperty(DataLakeGen2MetadataHandler.PARTITION_NUMBER)).thenReturn("0");
ValueSet valueSet = getSingleValueSet("varcharTest");
Constraints constraints = Mockito.mock(Constraints.class);
Mockito.when(constraints.getSummary()).thenReturn(new ImmutableMap.Builder<String, ValueSet>().put("testCol4", valueSet).build());
String expectedSql = "SELECT `testCol1`, `testCol2`, `testCol3`, `testCol4` FROM `testSchema`.`testTable` WHERE (`testCol4` = ?)";
PreparedStatement expectedPreparedStatement = Mockito.mock(PreparedStatement.class);
Mockito.when(this.connection.prepareStatement(Mockito.eq(expectedSql))).thenReturn(expectedPreparedStatement);
PreparedStatement preparedStatement = this.dataLakeGen2RecordHandler.buildSplitSql(this.connection, "testCatalogName", tableName, schema, constraints, split);
Assert.assertEquals(expectedPreparedStatement, preparedStatement);
Mockito.verify(preparedStatement, Mockito.times(1)).setString(1, "varcharTest");
}
use of com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet in project aws-athena-query-federation by awslabs.
the class DynamoDBMetadataHandlerTest method doGetSplitsQuery.
@Test
public void doGetSplitsQuery() throws Exception {
Map<String, ValueSet> constraintsMap = new HashMap<>();
EquatableValueSet.Builder valueSet = EquatableValueSet.newBuilder(allocator, Types.MinorType.VARCHAR.getType(), true, false);
for (int i = 0; i < 2000; i++) {
valueSet.add("test_str_" + i);
}
constraintsMap.put("col_0", valueSet.build());
GetTableLayoutResponse layoutResponse = 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));
GetSplitsRequest req = new GetSplitsRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, TEST_TABLE_NAME, layoutResponse.getPartitions(), ImmutableList.of("col_0"), new Constraints(new HashMap<>()), null);
logger.info("doGetSplits: req[{}]", req);
GetSplitsResponse response = handler.doGetSplits(allocator, req);
assertThat(response.getRequestType(), equalTo(MetadataRequestType.GET_SPLITS));
String continuationToken = response.getContinuationToken();
logger.info("doGetSplits: continuationToken[{}] - numSplits[{}]", continuationToken, response.getSplits().size());
assertThat(continuationToken, equalTo(String.valueOf(MAX_SPLITS_PER_REQUEST - 1)));
assertThat(response.getSplits().size(), equalTo(MAX_SPLITS_PER_REQUEST));
assertThat(response.getSplits().stream().map(split -> split.getProperty("col_0")).distinct().count(), equalTo((long) MAX_SPLITS_PER_REQUEST));
response = handler.doGetSplits(allocator, new GetSplitsRequest(req, continuationToken));
logger.info("doGetSplits: continuationToken[{}] - numSplits[{}]", continuationToken, response.getSplits().size());
assertThat(response.getContinuationToken(), equalTo(null));
assertThat(response.getSplits().size(), equalTo(MAX_SPLITS_PER_REQUEST));
assertThat(response.getSplits().stream().map(split -> split.getProperty("col_0")).distinct().count(), equalTo((long) MAX_SPLITS_PER_REQUEST));
}
Aggregations