private synchronized Map<PlanNodeId, TaskSource> updateSources(List<TaskSource> sources) {
Map<PlanNodeId, TaskSource> updatedUnpartitionedSources = new HashMap<>();
// don't update maxAcknowledgedSplit until the end because task sources may not
// be in sorted order and if we updated early we could skip splits
long newMaxAcknowledgedSplit = maxAcknowledgedSplit;
for (TaskSource source : sources) {
PlanNodeId sourceId = source.getPlanNodeId();
if (sourceId.equals(partitionedSourceId)) {
// partitioned split
ImmutableList.Builder<DriverSplitRunner> runners = ImmutableList.builder();
for (ScheduledSplit scheduledSplit : source.getSplits()) {
// only add a split if we have not already scheduled it
if (scheduledSplit.getSequenceId() > maxAcknowledgedSplit) {
// create a new driver for the split
runners.add(partitionedDriverFactory.createDriverRunner(scheduledSplit));
newMaxAcknowledgedSplit = max(scheduledSplit.getSequenceId(), newMaxAcknowledgedSplit);
}
}
enqueueDrivers(false, runners.build());
if (source.isNoMoreSplits()) {
partitionedDriverFactory.setNoMoreSplits();
}
} else {
// unpartitioned split
// update newMaxAcknowledgedSplit
for (ScheduledSplit scheduledSplit : source.getSplits()) {
newMaxAcknowledgedSplit = max(scheduledSplit.getSequenceId(), newMaxAcknowledgedSplit);
}
// create new source
TaskSource newSource;
TaskSource currentSource = unpartitionedSources.get(sourceId);
if (currentSource == null) {
newSource = source;
} else {
newSource = currentSource.update(source);
}
// only record new source if something changed
if (newSource != currentSource) {
unpartitionedSources.put(sourceId, newSource);
updatedUnpartitionedSources.put(sourceId, newSource);
}
}
}
maxAcknowledgedSplit = newMaxAcknowledgedSplit;
return updatedUnpartitionedSources;
}
private void processNewSource(TaskSource source) {
checkLockHeld("Lock must be held to call processNewSources");
// create new source
Set<ScheduledSplit> newSplits;
TaskSource currentSource = currentSources.get(source.getPlanNodeId());
if (currentSource == null) {
newSplits = source.getSplits();
currentSources.put(source.getPlanNodeId(), source);
} else {
// merge the current source and the specified source
TaskSource newSource = currentSource.update(source);
// if this is not a new source, just return
if (newSource == currentSource) {
return;
}
// find the new splits to add
newSplits = Sets.difference(newSource.getSplits(), currentSource.getSplits());
currentSources.put(source.getPlanNodeId(), newSource);
}
// add new splits
for (ScheduledSplit newSplit : newSplits) {
Split split = newSplit.getSplit();
SourceOperator sourceOperator = sourceOperators.get(source.getPlanNodeId());
if (sourceOperator != null) {
sourceOperator.addSplit(split);
}
}
// set no more splits
if (source.isNoMoreSplits()) {
sourceOperators.get(source.getPlanNodeId()).noMoreSplits();
}
}
@Override
public String getInfo() {
return (partitionedSplit == null) ? "" : partitionedSplit.getSplit().getInfo().toString();
}