@Override
public <A extends Writable> A getAggregatedValue(String name) {
AggregationStorageMetadata metadata = Configuration.get().getAggregationMetadata(name);
if (metadata == null) {
return super.getAggregatedValue(name);
} else {
AggregationStorage aggregationStorage = inboundAggregationStorages.get(name);
if (aggregationStorage == null) {
int numSplits = metadata.getNumSplits();
AggregationStorageWrapper storageWrapper = new AggregationStorageWrapper();
for (int i = 0; i < numSplits; ++i) {
AggregationStorageWrapper storageWrapperSplit =
getAggregatedValue(Configuration.get().getAggregationSplitName(name, i));
storageWrapper.aggregate(storageWrapperSplit);
}
aggregationStorage = storageWrapper.getAggregationStorage();
if (aggregationStorage == null) {
aggregationStorage = aggregationStorageFactory.createAggregationStorage(name, metadata);
}
inboundAggregationStorages.put(name, aggregationStorage);
}
return (A) aggregationStorage;
}
}
@Override
public void initAggregations() {
super.initAggregations();
Configuration.get()
.registerAggregation(
AGG_OUTPUT, IntWritable.class, LongWritable.class, true, new LongSumReduction());
}
private void splitAggregationStoragesAndSend(
String name, AggregationStorageMetadata metadata, AggregationStorage aggregationStorage) {
Configuration conf = Configuration.get();
int numSplits = metadata.getNumSplits();
if (numSplits == 1) {
AggregationStorageWrapper aggWrapper = new AggregationStorageWrapper(aggregationStorage);
super.setAggregatedValue(conf.getAggregationSplitName(name, 0), aggWrapper);
} else {
Collection<Writable> keysToTransfer = new ArrayList<>(aggregationStorage.getNumberMappings());
for (int i = 0; i < numSplits; ++i) {
AggregationStorage aggStorageSplit =
aggregationStorageFactory.createAggregationStorage(name);
AggregationStorageWrapper aggStorageSplitWrapper =
new AggregationStorageWrapper(aggStorageSplit);
keysToTransfer.clear();
for (Object key : aggregationStorage.getKeys()) {
int owner = key.hashCode() % numSplits;
if (owner < 0) {
owner += numSplits;
}
if (owner != i) {
continue;
}
keysToTransfer.add((Writable) key);
}
for (Writable key : keysToTransfer) {
aggStorageSplit.transferKeyFrom(key, aggregationStorage);
}
super.setAggregatedValue(conf.getAggregationSplitName(name, i), aggStorageSplitWrapper);
}
}
}
private void registerAggregationStore(
String aggName, AggregationStorageMetadata aggStorageMetadata)
throws IllegalAccessException, InstantiationException {
boolean isPersistent = aggStorageMetadata.isPersistent();
int numSplits = aggStorageMetadata.getNumSplits();
for (int i = 0; i < numSplits; ++i) {
String splitName = Configuration.get().getAggregationSplitName(aggName, i);
if (isPersistent) {
registerPersistentAggregator(splitName, AggregationStorageAggregator.class);
} else {
registerAggregator(splitName, AggregationStorageAggregator.class);
}
}
}
@Override
public <A extends Writable> void setAggregatedValue(String name, A value) {
AggregationStorageMetadata metadata = Configuration.get().getAggregationMetadata(name);
if (metadata == null) {
super.setAggregatedValue(name, value);
} else {
if (!(value instanceof AggregationStorage)) {
throw new RuntimeException(
"Value of an Arabesque aggregation should be a subclass of AggregationStorage");
}
AggregationStorage aggregationStorage = (AggregationStorage) value;
splitAggregationStoragesAndSend(name, metadata, aggregationStorage);
}
}
@Override
public void initialize() throws InstantiationException, IllegalAccessException {
// If we run with separate master, we need to initialize the worker context
// ourselves because a worker context is never created (and the worker context
// is the one that initializes our configuration according to the user settings).
WorkerContext workerContext = (WorkerContext) getConf().createWorkerContext();
workerContext.setConf(getConf());
Configuration conf = Configuration.get();
CommunicationStrategy communicationStrategy =
conf.createCommunicationStrategy(conf, null, workerContext);
numPhases = communicationStrategy.getNumPhases();
registeredAggregatorNames = new ArrayList<>();
savedAggregatorValues = new HashMap<>();
registerAggregator(AGG_EMBEDDINGS_GENERATED, LongSumAggregator.class);
registerAggregator(AGG_EMBEDDINGS_PROCESSED, LongSumAggregator.class);
registerAggregator(AGG_PROCESSED_SIZE_ODAG, LongMaxAggregator.class);
registerAggregator(AGG_PROCESSED_SIZE_CACHE, LongSumAggregator.class);
registerAggregator(AGG_CHILDREN_EVALUATED, LongSumAggregator.class);
registerAggregator(AGG_EMBEDDINGS_OUTPUT, LongSumAggregator.class);
Computation<?> computation = conf.createComputation();
computation.initAggregations();
Map<String, AggregationStorageMetadata> registeredAggregationStorages =
conf.getAggregationsMetadata();
LOG.info("Registered aggregation storages: " + registeredAggregationStorages);
for (Map.Entry<String, AggregationStorageMetadata> entry :
registeredAggregationStorages.entrySet()) {
String aggName = entry.getKey();
AggregationStorageMetadata aggStorageMetadata = entry.getValue();
registerAggregationStore(aggName, aggStorageMetadata);
}
masterComputation = conf.createMasterComputation();
masterComputation.setUnderlyingExecutionEngine(this);
masterComputation.init();
}