Search in sources :

Example 1 with FloatArrayElement

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);
}
Also used : ServerMatrix(com.tencent.angel.ps.storage.matrix.ServerMatrix) RowBasedPartition(com.tencent.angel.ps.storage.partition.RowBasedPartition) ServerLongAnyRow(com.tencent.angel.ps.storage.vector.ServerLongAnyRow) Random(java.util.Random) FloatArrayElement(com.tencent.angel.ps.storage.vector.element.FloatArrayElement) ServerPartition(com.tencent.angel.ps.storage.partition.ServerPartition)

Aggregations

ServerMatrix (com.tencent.angel.ps.storage.matrix.ServerMatrix)1 RowBasedPartition (com.tencent.angel.ps.storage.partition.RowBasedPartition)1 ServerPartition (com.tencent.angel.ps.storage.partition.ServerPartition)1 ServerLongAnyRow (com.tencent.angel.ps.storage.vector.ServerLongAnyRow)1 FloatArrayElement (com.tencent.angel.ps.storage.vector.element.FloatArrayElement)1 Random (java.util.Random)1