Examples with ILongKeyIntValuePartOp com.tencent.angel.psagent.matrix.transport.router.operator.ILongKeyIntValuePartOp used on opensource projects

Example 1 with ILongKeyIntValuePartOp

use of com.tencent.angel.psagent.matrix.transport.router.operator.ILongKeyIntValuePartOp in project angel by Tencent.

the class GetNeighborsWithCount method partitionGetWithCount.

public static PartitionGetResult partitionGetWithCount(PSContext psContext, PartitionGetParam partParam) {
    PartGetNeighborWithCountParam param = (PartGetNeighborWithCountParam) partParam;
    KeyValuePart keyValuePart = param.getIndicesPart();
    // Long type node id
    long[] nodeIds = ((ILongKeyIntValuePartOp) keyValuePart).getKeys();
    int[] count = ((ILongKeyIntValuePartOp) keyValuePart).getValues();
    long[][] neighbors = new long[nodeIds.length][];
    Random r = new Random();
    ServerLongAnyRow row = GraphMatrixUtils.getPSLongKeyRow(psContext, param);
    for (int i = 0; i < nodeIds.length; i++) {
        long nodeId = nodeIds[i];
        // Get node neighbor number
        NeighborsAliasTableElement element = (NeighborsAliasTableElement) (row.get(nodeId));
        if (element == null) {
            neighbors[i] = null;
        } else {
            long[] nodeNeighbors = element.getNeighborIds();
            if (nodeNeighbors == null || nodeNeighbors.length == 0 || count[i] <= 0) {
                neighbors[i] = null;
            } else {
                neighbors[i] = new long[count[i]];
                // start sampling by alias table for count times
                float[] accept = element.getAccept();
                int[] alias = element.getAlias();
                for (int j = 0; j < count[i]; j++) {
                    // int index = Math.abs(r.nextInt()) % nodeNeighbors.length;
                    int index = r.nextInt(nodeNeighbors.length);
                    // float ac = Math.abs(r.nextFloat());
                    float ac = r.nextFloat();
                    if (ac < accept[index]) {
                        neighbors[i][j] = nodeNeighbors[index];
                    } else {
                        neighbors[i][j] = nodeNeighbors[alias[index]];
                    }
                }
            }
        }
    }
    return new PartGetNeighborWithCountResult(nodeIds, neighbors);
}