@Override
public PlanNode visitExchange(
ExchangeNode exchange, RewriteContext<FragmentProperties> context) {
ImmutableList.Builder<SubPlan> builder = ImmutableList.builder();
if (exchange.getType() == ExchangeNode.Type.GATHER) {
context.get().setSingleNodeDistribution();
for (int i = 0; i < exchange.getSources().size(); i++) {
FragmentProperties childProperties = new FragmentProperties();
childProperties.setUnpartitionedOutput();
childProperties.setOutputLayout(exchange.getInputs().get(i));
builder.add(buildSubPlan(exchange.getSources().get(i), childProperties, context));
}
} else if (exchange.getType() == ExchangeNode.Type.REPARTITION) {
context.get().setFixedDistribution();
FragmentProperties childProperties =
new FragmentProperties()
.setPartitionedOutput(exchange.getPartitionKeys(), exchange.getHashSymbol())
.setOutputLayout(Iterables.getOnlyElement(exchange.getInputs()));
builder.add(
buildSubPlan(
Iterables.getOnlyElement(exchange.getSources()), childProperties, context));
} else if (exchange.getType() == ExchangeNode.Type.REPARTITION_WITH_NULL_REPLICATION) {
context.get().setFixedDistribution();
FragmentProperties childProperties =
new FragmentProperties()
.setPartitionedOutput(exchange.getPartitionKeys(), exchange.getHashSymbol())
.replicateNulls()
.setOutputLayout(Iterables.getOnlyElement(exchange.getInputs()));
builder.add(
buildSubPlan(
Iterables.getOnlyElement(exchange.getSources()), childProperties, context));
} else if (exchange.getType() == ExchangeNode.Type.REPLICATE) {
FragmentProperties childProperties = new FragmentProperties();
childProperties.setUnpartitionedOutput();
childProperties.setOutputLayout(Iterables.getOnlyElement(exchange.getInputs()));
builder.add(
buildSubPlan(
Iterables.getOnlyElement(exchange.getSources()), childProperties, context));
}
List<SubPlan> children = builder.build();
context.get().addChildren(children);
List<PlanFragmentId> childrenIds =
children
.stream()
.map(SubPlan::getFragment)
.map(PlanFragment::getId)
.collect(toImmutableList());
return new RemoteSourceNode(exchange.getId(), childrenIds, exchange.getOutputSymbols());
}
@Override
public PlanNode visitExchange(ExchangeNode node, RewriteContext<Set<Symbol>> context) {
Set<Symbol> expectedOutputSymbols = Sets.newHashSet(context.get());
node.getPartitionFunction().getHashColumn().ifPresent(expectedOutputSymbols::add);
node.getPartitionFunction()
.getPartitionFunctionArguments()
.stream()
.filter(PartitionFunctionArgumentBinding::isVariable)
.map(PartitionFunctionArgumentBinding::getColumn)
.forEach(expectedOutputSymbols::add);
List<List<Symbol>> inputsBySource = new ArrayList<>(node.getInputs().size());
for (int i = 0; i < node.getInputs().size(); i++) {
inputsBySource.add(new ArrayList<>());
}
List<Symbol> newOutputSymbols = new ArrayList<>(node.getOutputSymbols().size());
for (int i = 0; i < node.getOutputSymbols().size(); i++) {
Symbol outputSymbol = node.getOutputSymbols().get(i);
if (expectedOutputSymbols.contains(outputSymbol)) {
newOutputSymbols.add(outputSymbol);
for (int source = 0; source < node.getInputs().size(); source++) {
inputsBySource.get(source).add(node.getInputs().get(source).get(i));
}
}
}
// newOutputSymbols contains all partition and hash symbols so simply swap the output layout
PartitionFunctionBinding partitionFunctionBinding =
new PartitionFunctionBinding(
node.getPartitionFunction().getPartitioningHandle(),
newOutputSymbols,
node.getPartitionFunction().getPartitionFunctionArguments(),
node.getPartitionFunction().getHashColumn(),
node.getPartitionFunction().isReplicateNulls(),
node.getPartitionFunction().getBucketToPartition());
ImmutableList.Builder<PlanNode> rewrittenSources = ImmutableList.<PlanNode>builder();
for (int i = 0; i < node.getSources().size(); i++) {
ImmutableSet.Builder<Symbol> expectedInputs =
ImmutableSet.<Symbol>builder().addAll(inputsBySource.get(i));
rewrittenSources.add(context.rewrite(node.getSources().get(i), expectedInputs.build()));
}
return new ExchangeNode(
node.getId(),
node.getType(),
partitionFunctionBinding,
rewrittenSources.build(),
inputsBySource);
}