return planRewriter.rewrite(input, null);
}
}).list();
if (Iterables.all(rewrittenSources, not(instanceOf(MaterializeSampleNode.class)))) {
return new UnionNode(node.getId(), rewrittenSources, node.getSymbolMapping());
}
// Add sample weight to any sources that don't have it, and pull MaterializeSample through the UNION
ImmutableListMultimap.Builder<Symbol, Symbol> symbolMapping = ImmutableListMultimap.<Symbol, Symbol>builder().putAll(node.getSymbolMapping());
ImmutableList.Builder<PlanNode> sources = ImmutableList.builder();
Symbol outputSymbol = symbolAllocator.newSymbol("$sampleWeight", BIGINT);
for (PlanNode source : rewrittenSources) {
if (source instanceof MaterializeSampleNode) {
symbolMapping.put(outputSymbol, ((MaterializeSampleNode) source).getSampleWeightSymbol());
sources.add(((MaterializeSampleNode) source).getSource());
}
else {
Symbol symbol = symbolAllocator.newSymbol("$sampleWeight", BIGINT);
symbolMapping.put(outputSymbol, symbol);
sources.add(addSampleWeight(source, symbol));
}
}
node = new UnionNode(node.getId(), sources.build(), symbolMapping.build());
return new MaterializeSampleNode(idAllocator.getNextId(), node, outputSymbol);
}