public static Map<Symbol, Symbol> exchangeInputToOutput(ExchangeNode node, int sourceIndex) {
List<Symbol> inputSymbols = node.getInputs().get(sourceIndex);
Map<Symbol, Symbol> inputToOutput = new HashMap<>();
for (int i = 0; i < node.getOutputSymbols().size(); i++) {
inputToOutput.put(inputSymbols.get(i), node.getOutputSymbols().get(i));
}
return inputToOutput;
}
@Override
public Void visitExchange(ExchangeNode node, Void context) {
for (PlanFragmentId planFragmentId : node.getSourceFragmentIds()) {
PlanFragment target = fragmentsById.get(planFragmentId);
printEdge(node, target.getRoot());
}
return null;
}
@Override
public Void visitExchange(ExchangeNode node, Void context) {
List<Symbol> symbols = node.getOutputSymbols();
if (node.getType() == REPARTITION || node.getType() == REPARTITION_WITH_NULL_REPLICATION) {
symbols = node.getPartitionKeys().orElseGet(() -> ImmutableList.of(new Symbol("(absent)")));
}
String columns = Joiner.on(", ").join(symbols);
printNode(
node,
format("ExchangeNode[%s]", node.getType()),
columns,
NODE_COLORS.get(NodeType.EXCHANGE));
for (PlanNode planNode : node.getSources()) {
planNode.accept(this, context);
}
return null;
}
@Override
public ActualProperties visitExchange(
ExchangeNode node, List<ActualProperties> inputProperties) {
checkArgument(
node.getScope() != REMOTE
|| inputProperties.stream().noneMatch(ActualProperties::isNullsReplicated),
"Null replicated inputs should not be remotely exchanged");
Set<Map.Entry<Symbol, NullableValue>> entries = null;
for (int sourceIndex = 0; sourceIndex < node.getSources().size(); sourceIndex++) {
Map<Symbol, Symbol> inputToOutput = exchangeInputToOutput(node, sourceIndex);
ActualProperties translated =
inputProperties
.get(sourceIndex)
.translate(symbol -> Optional.ofNullable(inputToOutput.get(symbol)));
entries =
(entries == null)
? translated.getConstants().entrySet()
: Sets.intersection(entries, translated.getConstants().entrySet());
}
checkState(entries != null);
Map<Symbol, NullableValue> constants =
entries.stream().collect(toMap(Map.Entry::getKey, Map.Entry::getValue));
// Local exchanges are only created in AddLocalExchanges, at the end of optimization, and
// local exchanges do not produce global properties as represented by ActualProperties.
// This is acceptable because AddLocalExchanges does not use global properties and is only
// interested in the local properties.
// TODO: implement full properties for local exchanges
if (node.getScope() == LOCAL) {
return ActualProperties.builder().constants(constants).build();
}
switch (node.getType()) {
case GATHER:
boolean coordinatorOnly =
node.getPartitioningScheme().getPartitioning().getHandle().isCoordinatorOnly();
return ActualProperties.builder()
.global(
coordinatorOnly ? coordinatorSingleStreamPartition() : singleStreamPartition())
.constants(constants)
.build();
case REPARTITION:
return ActualProperties.builder()
.global(
partitionedOn(
node.getPartitioningScheme().getPartitioning(),
Optional.of(node.getPartitioningScheme().getPartitioning()))
.withReplicatedNulls(node.getPartitioningScheme().isReplicateNulls()))
.constants(constants)
.build();
case REPLICATE:
// TODO: this should have the same global properties as the stream taking the replicated
// data
return ActualProperties.builder()
.global(arbitraryPartition())
.constants(constants)
.build();
}
throw new UnsupportedOperationException("not yet implemented");
}
@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);
}