@Override
public Collection<Partition<T>> definePartitions(
Collection<Partition<T>> partitions, PartitioningContext context) {
if (partitionedInstanceStatus == null || partitionedInstanceStatus.isEmpty()) {
// first call
// trying to give initial stability before sending the repartition call
if (partitionedInstanceStatus == null) {
partitionedInstanceStatus = new HashMap<Integer, BatchedOperatorStats>();
}
partitionNextMillis = System.currentTimeMillis() + 2 * cooldownMillis;
nextMillis = partitionNextMillis;
// delegate to create initial list of partitions
return new StatelessPartitioner<T>(initialPartitionCount)
.definePartitions(partitions, context);
} else {
// repartition call
logger.debug("repartition call for operator");
if (System.currentTimeMillis() < partitionNextMillis) {
return partitions;
}
List<Partition<T>> newPartitions = new ArrayList<Partition<T>>();
HashMap<Integer, Partition<T>> lowLoadPartitions = new HashMap<Integer, Partition<T>>();
for (Partition<T> p : partitions) {
int load = getLoad(p.getStats());
if (load < 0) {
// combine neighboring underutilized partitions
PartitionKeys pks =
p.getPartitionKeys().values().iterator().next(); // one port partitioned
for (int partitionKey : pks.partitions) {
// look for the sibling partition by excluding leading bit
int reducedMask = pks.mask >>> 1;
String lookupKey =
Integer.valueOf(reducedMask) + "-" + Integer.valueOf(partitionKey & reducedMask);
logger.debug("pks {} lookupKey {}", pks, lookupKey);
Partition<T> siblingPartition = lowLoadPartitions.remove(partitionKey & reducedMask);
if (siblingPartition == null) {
lowLoadPartitions.put(partitionKey & reducedMask, p);
} else {
// both of the partitions are low load, combine
PartitionKeys newPks =
new PartitionKeys(reducedMask, Sets.newHashSet(partitionKey & reducedMask));
// put new value so the map gets marked as modified
Operator.InputPort<?> port =
siblingPartition.getPartitionKeys().keySet().iterator().next();
siblingPartition.getPartitionKeys().put(port, newPks);
// add as new partition
newPartitions.add(siblingPartition);
// LOG.debug("partition keys after merge {}", siblingPartition.getPartitionKeys());
}
}
} else if (load > 0) {
// split bottlenecks
Map<Operator.InputPort<?>, PartitionKeys> keys = p.getPartitionKeys();
Map.Entry<Operator.InputPort<?>, PartitionKeys> e = keys.entrySet().iterator().next();
final int newMask;
final Set<Integer> newKeys;
if (e.getValue().partitions.size() == 1) {
// split single key
newMask = (e.getValue().mask << 1) | 1;
int key = e.getValue().partitions.iterator().next();
int key2 = (newMask ^ e.getValue().mask) | key;
newKeys = Sets.newHashSet(key, key2);
} else {
// assign keys to separate partitions
newMask = e.getValue().mask;
newKeys = e.getValue().partitions;
}
for (int key : newKeys) {
Partition<T> newPartition = new DefaultPartition<T>(p.getPartitionedInstance());
newPartition
.getPartitionKeys()
.put(e.getKey(), new PartitionKeys(newMask, Sets.newHashSet(key)));
newPartitions.add(newPartition);
}
} else {
// leave unchanged
newPartitions.add(p);
}
}
// put back low load partitions that could not be combined
newPartitions.addAll(lowLoadPartitions.values());
partitionNextMillis = System.currentTimeMillis() + cooldownMillis;
return newPartitions;
}
}
