use of com.tencent.angel.graph.model.neighbor.simplewithtype.TypeNeighborElement in project angel by Tencent.
the class Sample method partitionGet.
@Override
public PartitionGetResult partitionGet(PartitionGetParam partParam) {
PartSampleParam param = (PartSampleParam) partParam;
KeyPart keyPart = param.getIndicesPart();
int sampleType = param.getSampleType();
long[] nodeIds = ((ILongKeyPartOp) keyPart).getKeys();
Long2IntOpenHashMap nodeIdToSizes = new Long2IntOpenHashMap(nodeIds.length);
for (int i = 0; i < nodeIds.length; i++) {
nodeIdToSizes.addTo(nodeIds[i], 1);
}
ServerLongAnyRow row = GraphMatrixUtils.getPSLongKeyRow(psContext, param);
Random r = new Random();
long[] distinctNodeIds = new long[nodeIdToSizes.size()];
long[][] samples = new long[nodeIdToSizes.size()][];
ObjectIterator<Entry> iter = nodeIdToSizes.long2IntEntrySet().fastIterator();
int index = 0;
while (iter.hasNext()) {
Entry entry = iter.next();
distinctNodeIds[index] = entry.getLongKey();
TypeNeighborElement element = (TypeNeighborElement) row.get(distinctNodeIds[index]);
samples[index] = element.Sample(sampleType, r, distinctNodeIds[index], entry.getIntValue());
index++;
}
return new PartGetLongsResult(distinctNodeIds, samples);
}
Aggregations