if (rightSource instanceof MaterializeSampleNode) {
rightSampleWeight = ((MaterializeSampleNode) rightSource).getSampleWeightSymbol();
rightSource = ((MaterializeSampleNode) rightSource).getSource();
}
PlanNode joinNode = new JoinNode(node.getId(), node.getType(), leftSource, rightSource, node.getCriteria());
Symbol outputSampleWeight;
if (leftSampleWeight != null && rightSampleWeight != null) {
ImmutableMap.Builder<Symbol, Expression> projections = ImmutableMap.builder();
Expression sampleWeightExpr;
switch (node.getType()) {
case INNER:
case CROSS:
sampleWeightExpr = new ArithmeticExpression(ArithmeticExpression.Type.MULTIPLY, new QualifiedNameReference(leftSampleWeight.toQualifiedName()), new QualifiedNameReference(rightSampleWeight.toQualifiedName()));
break;
case LEFT:
sampleWeightExpr = new ArithmeticExpression(ArithmeticExpression.Type.MULTIPLY, new QualifiedNameReference(leftSampleWeight.toQualifiedName()), oneIfNull(rightSampleWeight));
break;
case RIGHT:
sampleWeightExpr = new ArithmeticExpression(ArithmeticExpression.Type.MULTIPLY, oneIfNull(leftSampleWeight), new QualifiedNameReference(rightSampleWeight.toQualifiedName()));
break;
default:
throw new AssertionError(String.format("Unknown join type: %s", node.getType()));
}
outputSampleWeight = symbolAllocator.newSymbol(sampleWeightExpr, BIGINT);
projections.put(outputSampleWeight, sampleWeightExpr);
for (Symbol symbol : Iterables.filter(node.getOutputSymbols(), not(in(ImmutableSet.of(leftSampleWeight, rightSampleWeight))))) {
Expression expression = new QualifiedNameReference(symbol.toQualifiedName());
projections.put(symbol, expression);
}
joinNode = new ProjectNode(idAllocator.getNextId(), joinNode, projections.build());
}
else {
outputSampleWeight = leftSampleWeight == null ? rightSampleWeight : leftSampleWeight;
if ((node.getType() == JoinNode.Type.LEFT && leftSampleWeight == null) || (node.getType() == JoinNode.Type.RIGHT && rightSampleWeight == null)) {
// There could be NULLs in the sample weight, so fix them with a projection
ImmutableMap.Builder<Symbol, Expression> projections = ImmutableMap.builder();
for (Symbol symbol : Iterables.filter(node.getOutputSymbols(), not(equalTo(outputSampleWeight)))) {
Expression expression = new QualifiedNameReference(symbol.toQualifiedName());
projections.put(symbol, expression);
}
Expression sampleWeightExpr = oneIfNull(outputSampleWeight);
outputSampleWeight = symbolAllocator.newSymbol(sampleWeightExpr, BIGINT);
projections.put(outputSampleWeight, sampleWeightExpr);
joinNode = new ProjectNode(idAllocator.getNextId(), joinNode, projections.build());
}
}
return new MaterializeSampleNode(idAllocator.getNextId(), joinNode, outputSampleWeight);
}
return new JoinNode(node.getId(), node.getType(), leftSource, rightSource, node.getCriteria());
}