private SubPlanBuilder addDistributedAggregation(
SubPlanBuilder plan,
Map<Symbol, FunctionCall> aggregations,
Map<Symbol, Signature> functions,
Map<Symbol, Symbol> masks,
List<Symbol> groupBy,
Optional<Symbol> sampleWeight,
double confidence) {
Map<Symbol, FunctionCall> finalCalls = new HashMap<>();
Map<Symbol, FunctionCall> intermediateCalls = new HashMap<>();
Map<Symbol, Signature> intermediateFunctions = new HashMap<>();
Map<Symbol, Symbol> intermediateMask = new HashMap<>();
for (Map.Entry<Symbol, FunctionCall> entry : aggregations.entrySet()) {
Signature signature = functions.get(entry.getKey());
FunctionInfo function = metadata.getFunction(signature);
Symbol intermediateSymbol =
allocator.newSymbol(function.getName().getSuffix(), function.getIntermediateType());
intermediateCalls.put(intermediateSymbol, entry.getValue());
intermediateFunctions.put(intermediateSymbol, signature);
if (masks.containsKey(entry.getKey())) {
intermediateMask.put(intermediateSymbol, masks.get(entry.getKey()));
}
// rewrite final aggregation in terms of intermediate function
finalCalls.put(
entry.getKey(),
new FunctionCall(
function.getName(),
ImmutableList.<Expression>of(
new QualifiedNameReference(intermediateSymbol.toQualifiedName()))));
}
// create partial aggregation plan
AggregationNode partialAggregation =
new AggregationNode(
idAllocator.getNextId(),
plan.getRoot(),
groupBy,
intermediateCalls,
intermediateFunctions,
intermediateMask,
PARTIAL,
sampleWeight,
confidence);
plan.setRoot(
new SinkNode(
idAllocator.getNextId(), partialAggregation, partialAggregation.getOutputSymbols()));
// create final aggregation plan
ExchangeNode source =
new ExchangeNode(
idAllocator.getNextId(), plan.getId(), plan.getRoot().getOutputSymbols());
AggregationNode finalAggregation =
new AggregationNode(
idAllocator.getNextId(),
source,
groupBy,
finalCalls,
functions,
ImmutableMap.<Symbol, Symbol>of(),
FINAL,
Optional.<Symbol>absent(),
confidence);
if (groupBy.isEmpty()) {
plan = createSingleNodePlan(finalAggregation).addChild(plan.build());
} else {
plan.setHashOutputPartitioning(groupBy);
plan = createFixedDistributionPlan(finalAggregation).addChild(plan.build());
}
return plan;
}
