@Override
public SubPlanBuilder visitAggregation(AggregationNode node, Void context) {
SubPlanBuilder current = node.getSource().accept(this, context);
if (!current.isDistributed()) {
// add the aggregation node as the root of the current fragment
current.setRoot(
new AggregationNode(
node.getId(),
current.getRoot(),
node.getGroupBy(),
node.getAggregations(),
node.getFunctions(),
node.getMasks(),
SINGLE,
node.getSampleWeight(),
node.getConfidence()));
return current;
}
Map<Symbol, FunctionCall> aggregations = node.getAggregations();
Map<Symbol, Signature> functions = node.getFunctions();
Map<Symbol, Symbol> masks = node.getMasks();
List<Symbol> groupBy = node.getGroupBy();
boolean decomposable = true;
for (Signature function : functions.values()) {
if (!metadata.getFunction(function).getAggregationFunction().isDecomposable()) {
decomposable = false;
break;
}
}
// else, we need to "close" the current fragment and create an unpartitioned fragment for the
// final aggregation
if (decomposable) {
return addDistributedAggregation(
current,
aggregations,
functions,
masks,
groupBy,
node.getSampleWeight(),
node.getConfidence());
}
return addSingleNodeAggregation(
current,
aggregations,
functions,
masks,
groupBy,
node.getSampleWeight(),
node.getConfidence());
}
@Override
public PlanNode rewriteAggregation(
AggregationNode node,
Expression inheritedPredicate,
PlanRewriter<Expression> planRewriter) {
EqualityInference equalityInference = createEqualityInference(inheritedPredicate);
List<Expression> pushdownConjuncts = new ArrayList<>();
List<Expression> postAggregationConjuncts = new ArrayList<>();
// Strip out non-deterministic conjuncts
postAggregationConjuncts.addAll(
ImmutableList.copyOf(filter(extractConjuncts(inheritedPredicate), not(deterministic()))));
inheritedPredicate = stripNonDeterministicConjuncts(inheritedPredicate);
// Sort non-equality predicates by those that can be pushed down and those that cannot
for (Expression conjunct : EqualityInference.nonInferrableConjuncts(inheritedPredicate)) {
Expression rewrittenConjunct =
equalityInference.rewriteExpression(conjunct, in(node.getGroupBy()));
if (rewrittenConjunct != null) {
pushdownConjuncts.add(rewrittenConjunct);
} else {
postAggregationConjuncts.add(conjunct);
}
}
// Add the equality predicates back in
EqualityInference.EqualityPartition equalityPartition =
equalityInference.generateEqualitiesPartitionedBy(in(node.getGroupBy()));
pushdownConjuncts.addAll(equalityPartition.getScopeEqualities());
postAggregationConjuncts.addAll(equalityPartition.getScopeComplementEqualities());
postAggregationConjuncts.addAll(equalityPartition.getScopeStraddlingEqualities());
PlanNode rewrittenSource =
planRewriter.rewrite(node.getSource(), combineConjuncts(pushdownConjuncts));
PlanNode output = node;
if (rewrittenSource != node.getSource()) {
output =
new AggregationNode(
node.getId(),
rewrittenSource,
node.getGroupBy(),
node.getAggregations(),
node.getFunctions(),
node.getMasks(),
node.getStep(),
node.getSampleWeight(),
node.getConfidence());
}
if (!postAggregationConjuncts.isEmpty()) {
output =
new FilterNode(
idAllocator.getNextId(), output, combineConjuncts(postAggregationConjuncts));
}
return output;
}
@Override
public PlanNode visitAggregation(AggregationNode node, RewriteContext<Set<Symbol>> context) {
ImmutableSet.Builder<Symbol> expectedInputs =
ImmutableSet.<Symbol>builder().addAll(node.getGroupBy());
if (node.getHashSymbol().isPresent()) {
expectedInputs.add(node.getHashSymbol().get());
}
ImmutableMap.Builder<Symbol, Signature> functions = ImmutableMap.builder();
ImmutableMap.Builder<Symbol, FunctionCall> functionCalls = ImmutableMap.builder();
ImmutableMap.Builder<Symbol, Symbol> masks = ImmutableMap.builder();
for (Map.Entry<Symbol, FunctionCall> entry : node.getAggregations().entrySet()) {
Symbol symbol = entry.getKey();
if (context.get().contains(symbol)) {
FunctionCall call = entry.getValue();
expectedInputs.addAll(DependencyExtractor.extractUnique(call));
if (node.getMasks().containsKey(symbol)) {
expectedInputs.add(node.getMasks().get(symbol));
masks.put(symbol, node.getMasks().get(symbol));
}
functionCalls.put(symbol, call);
functions.put(symbol, node.getFunctions().get(symbol));
}
}
if (node.getSampleWeight().isPresent()) {
expectedInputs.add(node.getSampleWeight().get());
}
PlanNode source = context.rewrite(node.getSource(), expectedInputs.build());
return new AggregationNode(
node.getId(),
source,
node.getGroupBy(),
functionCalls.build(),
functions.build(),
masks.build(),
node.getStep(),
node.getSampleWeight(),
node.getConfidence(),
node.getHashSymbol());
}