use of com.tencent.angel.ps.storage.vector.element.FloatArrayElement in project angel by Tencent.
the class GetNodeAttrs method partitionGet.
@Override
public PartitionGetResult partitionGet(PartitionGetParam partParam) {
PartGetNodeAttrsParam param = (PartGetNodeAttrsParam) partParam;
ServerMatrix matrix = psContext.getMatrixStorageManager().getMatrix(partParam.getMatrixId());
ServerPartition part = matrix.getPartition(partParam.getPartKey().getPartitionId());
ServerLongAnyRow row = (ServerLongAnyRow) (((RowBasedPartition) part).getRow(0));
long[] nodeIds = param.getNodeIds();
float[][] attrs = new float[nodeIds.length][];
int count = param.getCount();
Random r = new Random();
for (int i = 0; i < nodeIds.length; i++) {
long nodeId = nodeIds[i];
// Get node neighbor number
FloatArrayElement element = (FloatArrayElement) (row.get(nodeId));
if (element == null) {
attrs[i] = null;
} else {
float[] nodeAttrs = element.getData();
if (nodeAttrs == null || nodeAttrs.length == 0) {
attrs[i] = null;
} else if (count <= 0 || nodeAttrs.length <= count) {
attrs[i] = nodeAttrs;
} else {
attrs[i] = new float[count];
// If the neighbor number > count, just copy a range of neighbors to the result array, the copy position is random
int startPos = Math.abs(r.nextInt()) % nodeAttrs.length;
if (startPos + count <= nodeAttrs.length) {
System.arraycopy(nodeAttrs, startPos, attrs[i], 0, count);
} else {
System.arraycopy(nodeAttrs, startPos, attrs[i], 0, nodeAttrs.length - startPos);
System.arraycopy(nodeAttrs, 0, attrs[i], nodeAttrs.length - startPos, count - (nodeAttrs.length - startPos));
}
}
}
}
return new PartGetNodeAttrsResult(part.getPartitionKey().getPartitionId(), attrs);
}
Aggregations