@Override
protected RelationPlan visitUnion(Union node, Void context) {
checkArgument(!node.getRelations().isEmpty(), "No relations specified for UNION");
List<Symbol> unionOutputSymbols = null;
ImmutableList.Builder<PlanNode> sources = ImmutableList.builder();
ImmutableListMultimap.Builder<Symbol, Symbol> symbolMapping = ImmutableListMultimap.builder();
List<RelationPlan> subPlans =
node.getRelations()
.stream()
.map(relation -> processAndCoerceIfNecessary(relation, context))
.collect(toImmutableList());
boolean hasSampleWeight = false;
for (RelationPlan subPlan : subPlans) {
if (subPlan.getSampleWeight().isPresent()) {
hasSampleWeight = true;
break;
}
}
Optional<Symbol> outputSampleWeight = Optional.empty();
for (RelationPlan relationPlan : subPlans) {
if (hasSampleWeight && !relationPlan.getSampleWeight().isPresent()) {
relationPlan = addConstantSampleWeight(relationPlan);
}
List<Symbol> childOutputSymbols = relationPlan.getOutputSymbols();
if (unionOutputSymbols == null) {
// Use the first Relation to derive output symbol names
TupleDescriptor descriptor = relationPlan.getDescriptor();
ImmutableList.Builder<Symbol> outputSymbolBuilder = ImmutableList.builder();
for (Field field : descriptor.getVisibleFields()) {
int fieldIndex = descriptor.indexOf(field);
Symbol symbol = childOutputSymbols.get(fieldIndex);
outputSymbolBuilder.add(
symbolAllocator.newSymbol(symbol.getName(), symbolAllocator.getTypes().get(symbol)));
}
unionOutputSymbols = outputSymbolBuilder.build();
outputSampleWeight = relationPlan.getSampleWeight();
}
TupleDescriptor descriptor = relationPlan.getDescriptor();
checkArgument(
descriptor.getVisibleFieldCount() == unionOutputSymbols.size(),
"Expected relation to have %s symbols but has %s symbols",
descriptor.getVisibleFieldCount(),
unionOutputSymbols.size());
int unionFieldId = 0;
for (Field field : descriptor.getVisibleFields()) {
int fieldIndex = descriptor.indexOf(field);
symbolMapping.put(unionOutputSymbols.get(unionFieldId), childOutputSymbols.get(fieldIndex));
unionFieldId++;
}
sources.add(relationPlan.getRoot());
}
PlanNode planNode =
new UnionNode(idAllocator.getNextId(), sources.build(), symbolMapping.build());
if (node.isDistinct()) {
planNode = distinct(planNode);
}
return new RelationPlan(
planNode,
analysis.getOutputDescriptor(node),
planNode.getOutputSymbols(),
outputSampleWeight);
}