@Override public PostAggregation replacePostAggWithPostAggFromMap( PostAggregation postAggregation, Map<String, Aggregation> oldNameToNewAggregationMapping ) { if (postAggregation instanceof FieldAccessorPostAggregation) { //Check if the aggregation which this postAgg is referencing in the outerQuery has changed. //if so, create new FieldAccessor that accesses the changed aggregation String fieldName = ((FieldAccessorPostAggregation) postAggregation).getFieldName(); if (oldNameToNewAggregationMapping.containsKey(fieldName)) { return new FieldAccessorPostAggregation(oldNameToNewAggregationMapping.get(fieldName)); } else { //The agg which this fieldAccessor is referencing has not changed. So return the fieldAccessor as it is. return postAggregation; } } else if (postAggregation instanceof WithFields) { List<PostAggregation> resultPostAggsList = new ArrayList<>(); @SuppressWarnings("unchecked") List<PostAggregation> childPostAggs = ((WithFields) postAggregation).getFields(); for (PostAggregation postAgg : childPostAggs) { resultPostAggsList.add(replacePostAggWithPostAggFromMap(postAgg, oldNameToNewAggregationMapping)); } return ((WithFields) postAggregation).withFields(resultPostAggsList); } else { return postAggregation; } }
@Override public TemplateDruidQuery updateOuterQuery( TemplateDruidQuery outerQuery, Map<String, String> oldFieldNameToNewFieldNameMap ) { Map<String, Aggregation> oldNameToNewAggregationMapping = new HashMap<>(); Set<Aggregation> updatedOuterAggs = updateQueryAggs( outerQuery.getAggregations(), oldFieldNameToNewFieldNameMap, oldNameToNewAggregationMapping ); //Update the FieldAccessors from the outer query post aggs to access the correct aggs. Set<PostAggregation> updateOuterPostAggs = new HashSet<>(); for (PostAggregation postAggregation: outerQuery.getPostAggregations()) { updateOuterPostAggs.add(replacePostAggWithPostAggFromMap(postAggregation, oldNameToNewAggregationMapping)); } //create new TDQ using updated aggs, updatedPostAggs, updatedInnerQuery and timegrain of outerQuery return new TemplateDruidQuery( updatedOuterAggs, updateOuterPostAggs, outerQuery.getInnerQuery().orElse(null), outerQuery.getTimeGrain() ); }