Source:
public DruidQueryGeneratorContext withAggregation( Map<VariableReferenceExpression, Selection> newSelections, Set<VariableReferenceExpression> newGroupByColumns, int newAggregations, Set<VariableReferenceExpression> newHiddenColumnSet) { AtomicBoolean pushDownDistinctCount = new AtomicBoolean(false); newSelections.values().forEach(selection -> { if (selection.getDefinition().startsWith(DRUID_COUNT_DISTINCT_FUNCTION_NAME.toUpperCase(Locale.ENGLISH))) { pushDownDistinctCount.set(true); } }); Map<VariableReferenceExpression, Selection> targetSelections = newSelections; if (pushDownDistinctCount.get()) { newHiddenColumnSet = ImmutableSet.of(); ImmutableMap.Builder<VariableReferenceExpression, Selection> builder = ImmutableMap.builder(); for (Map.Entry<VariableReferenceExpression, Selection> entry : newSelections.entrySet()) { if (entry.getValue().getDefinition().startsWith(DRUID_COUNT_DISTINCT_FUNCTION_NAME.toUpperCase(Locale.ENGLISH))) { String definition = entry.getValue().getDefinition(); int start = definition.indexOf("("); int end = definition.indexOf(")"); String countDistinctClause = "count ( distinct \"" + definition.substring(start + 1, end) + "\")"; Selection countDistinctSelection = new Selection(countDistinctClause, entry.getValue().getOrigin()); builder.put(entry.getKey(), countDistinctSelection); } else { builder.put(entry.getKey(), entry.getValue()); } } targetSelections = builder.build(); } else { checkState(!hasAggregation(), "Druid doesn't support aggregation on top of the aggregated data"); } checkState(!hasLimit(), "Druid doesn't support aggregation on top of the limit"); checkState(newAggregations > 0, "Invalid number of aggregations"); return new DruidQueryGeneratorContext( targetSelections, from, filter, limit, newAggregations, newGroupByColumns, variablesInAggregation, newHiddenColumnSet); }

Target:
public DruidQueryGeneratorContext withAggregation( Map<VariableReferenceExpression, Selection> newSelections, Set<VariableReferenceExpression> newGroupByColumns, int newAggregations, Set<VariableReferenceExpression> newHiddenColumnSet) { AtomicBoolean pushDownDistinctCount = new AtomicBoolean(false); newSelections.values().forEach(selection -> { if (selection.getDefinition().startsWith(DRUID_COUNT_DISTINCT_FUNCTION_NAME.toUpperCase(Locale.ENGLISH))) { pushDownDistinctCount.set(true); } }); Map<VariableReferenceExpression, Selection> targetSelections = newSelections; if (pushDownDistinctCount.get()) { newHiddenColumnSet = ImmutableSet.of(); ImmutableMap.Builder<VariableReferenceExpression, Selection> builder = ImmutableMap.builder(); for (Map.Entry<VariableReferenceExpression, Selection> entry : newSelections.entrySet()) { if (entry.getValue().getDefinition().startsWith(DRUID_COUNT_DISTINCT_FUNCTION_NAME.toUpperCase(Locale.ENGLISH))) { String definition = entry.getValue().getDefinition(); int start = definition.indexOf("("); int end = definition.indexOf(")"); String countDistinctClause = "count ( distinct " + escapeSqlIdentifier(definition.substring(start + 1, end)) + ")"; Selection countDistinctSelection = new Selection(countDistinctClause, entry.getValue().getOrigin()); builder.put(entry.getKey(), countDistinctSelection); } else { builder.put(entry.getKey(), entry.getValue()); } } targetSelections = builder.build(); } else { checkState(!hasAggregation(), "Druid doesn't support aggregation on top of the aggregated data"); } checkState(!hasLimit(), "Druid doesn't support aggregation on top of the limit"); checkState(newAggregations > 0, "Invalid number of aggregations"); return new DruidQueryGeneratorContext( targetSelections, from, filter, limit, newAggregations, newGroupByColumns, variablesInAggregation, newHiddenColumnSet); }