Examples with SEGMENT_COUNT_METADATA com.amazonaws.athena.connectors.dynamodb.constants.DynamoDBConstants.SEGMENT_COUNT_METADATA used on opensource projects

Example 1 with SEGMENT_COUNT_METADATA

use of com.amazonaws.athena.connectors.dynamodb.constants.DynamoDBConstants.SEGMENT_COUNT_METADATA in project aws-athena-query-federation by awslabs.

the class DynamoDBRecordHandler method buildReadRequest.

/*
    Converts a split into a Query or Scan request
     */
private AmazonWebServiceRequest buildReadRequest(Split split, String tableName, Schema schema) {
    validateExpectedMetadata(split.getProperties());
    // prepare filters
    String rangeKeyFilter = split.getProperty(RANGE_KEY_FILTER_METADATA);
    String nonKeyFilter = split.getProperty(NON_KEY_FILTER_METADATA);
    Map<String, String> expressionAttributeNames = new HashMap<>();
    Map<String, AttributeValue> expressionAttributeValues = new HashMap<>();
    if (rangeKeyFilter != null || nonKeyFilter != null) {
        try {
            expressionAttributeNames.putAll(Jackson.getObjectMapper().readValue(split.getProperty(EXPRESSION_NAMES_METADATA), STRING_MAP_TYPE_REFERENCE));
            expressionAttributeValues.putAll(Jackson.getObjectMapper().readValue(split.getProperty(EXPRESSION_VALUES_METADATA), ATTRIBUTE_VALUE_MAP_TYPE_REFERENCE));
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }
    // Only read columns that are needed in the query
    String projectionExpression = schema.getFields().stream().map(field -> {
        String aliasedName = DDBPredicateUtils.aliasColumn(field.getName());
        expressionAttributeNames.put(aliasedName, field.getName());
        return aliasedName;
    }).collect(Collectors.joining(","));
    boolean isQuery = split.getProperty(SEGMENT_ID_PROPERTY) == null;
    if (isQuery) {
        // prepare key condition expression
        String indexName = split.getProperty(INDEX_METADATA);
        String hashKeyName = split.getProperty(HASH_KEY_NAME_METADATA);
        String hashKeyAlias = DDBPredicateUtils.aliasColumn(hashKeyName);
        String keyConditionExpression = hashKeyAlias + " = " + HASH_KEY_VALUE_ALIAS;
        if (rangeKeyFilter != null) {
            keyConditionExpression += " AND " + rangeKeyFilter;
        }
        expressionAttributeNames.put(hashKeyAlias, hashKeyName);
        expressionAttributeValues.put(HASH_KEY_VALUE_ALIAS, Jackson.fromJsonString(split.getProperty(hashKeyName), AttributeValue.class));
        return new QueryRequest().withTableName(tableName).withIndexName(indexName).withKeyConditionExpression(keyConditionExpression).withFilterExpression(nonKeyFilter).withExpressionAttributeNames(expressionAttributeNames).withExpressionAttributeValues(expressionAttributeValues).withProjectionExpression(projectionExpression);
    } else {
        int segmentId = Integer.parseInt(split.getProperty(SEGMENT_ID_PROPERTY));
        int segmentCount = Integer.parseInt(split.getProperty(SEGMENT_COUNT_METADATA));
        return new ScanRequest().withTableName(tableName).withSegment(segmentId).withTotalSegments(segmentCount).withFilterExpression(nonKeyFilter).withExpressionAttributeNames(expressionAttributeNames.isEmpty() ? null : expressionAttributeNames).withExpressionAttributeValues(expressionAttributeValues.isEmpty() ? null : expressionAttributeValues).withProjectionExpression(projectionExpression);
    }
}