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

Example 21 with KeyPart

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

the class SampleEdgeFeatParam method split.

@Override
public List<PartitionGetParam> split() {
    // Get matrix meta
    MatrixMeta meta = PSAgentContext.get().getMatrixMetaManager().getMatrixMeta(matrixId);
    PartitionKey[] partitions = meta.getPartitionKeys();
    // Split nodeIds
    KeyPart[] splits = RouterUtils.split(meta, 0, nodeIds);
    assert partitions.length == splits.length;
    // Generate rpc params
    List<PartitionGetParam> partParams = new ArrayList<>(partitions.length);
    for (int i = 0; i < partitions.length; i++) {
        if (splits[i] != null && splits[i].size() > 0) {
            partParams.add(new PartSampleEdgeFeatParam(matrixId, partitions[i], splits[i], count));
        }
    }
    return partParams;
}

Example 22 with KeyPart

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

the class GetNodeAttrsParam method split.

@Override
public List<PartitionGetParam> split() {
    // Get matrix meta
    MatrixMeta meta = PSAgentContext.get().getMatrixMetaManager().getMatrixMeta(matrixId);
    PartitionKey[] parts = meta.getPartitionKeys();
    // Split
    KeyPart[] nodeIdsParts = RouterUtils.split(meta, 0, nodeIds, false);
    // Generate Part psf get param
    List<PartitionGetParam> partParams = new ArrayList<>(parts.length);
    assert parts.length == nodeIdsParts.length;
    for (int i = 0; i < parts.length; i++) {
        if (nodeIdsParts[i] != null && nodeIdsParts[i].size() > 0) {
            partParams.add(new GeneralPartGetParam(matrixId, parts[i], nodeIdsParts[i]));
        }
    }
    return partParams;
}

Example 23 with KeyPart

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

the class GetLabels method partitionGet.

@Override
public PartitionGetResult partitionGet(PartitionGetParam partParam) {
    GeneralPartGetParam param = (GeneralPartGetParam) partParam;
    KeyPart keyPart = param.getIndicesPart();
    switch(keyPart.getKeyType()) {
        case LONG:
            {
                // Long type node id
                long[] nodeIds = ((ILongKeyPartOp) keyPart).getKeys();
                ServerLongAnyRow row = GraphMatrixUtils.getPSLongKeyRow(psContext, param);
                Long2ObjectOpenHashMap<float[]> nodeIdToLabels = new Long2ObjectOpenHashMap<>(nodeIds.length);
                for (long nodeId : nodeIds) {
                    if (row.get(nodeId) == null) {
                        // If node not exist, just skip
                        continue;
                    }
                    float[] labels = ((GraphNode) (row.get(nodeId))).getLabels();
                    if (labels != null) {
                        nodeIdToLabels.put(nodeId, labels);
                    }
                }
                return new PartGetFloatArrayAttrsResult(param.getPartKey().getPartitionId(), nodeIdToLabels);
            }
        default:
            {
                // TODO: support String, Int, and Any type node id
                throw new InvalidParameterException("Unsupport index type " + keyPart.getKeyType());
            }
    }
}

Example 24 with KeyPart

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

the class GetNodeFeats method partitionGet.

@Override
public PartitionGetResult partitionGet(PartitionGetParam partParam) {
    GeneralPartGetParam param = (GeneralPartGetParam) partParam;
    KeyPart keyPart = param.getIndicesPart();
    switch(keyPart.getKeyType()) {
        case LONG:
            {
                // Long type node id
                long[] nodeIds = ((ILongKeyPartOp) keyPart).getKeys();
                ServerLongAnyRow row = GraphMatrixUtils.getPSLongKeyRow(psContext, param);
                Long2ObjectOpenHashMap<IntFloatVector> nodeIdToFeats = new Long2ObjectOpenHashMap<>(nodeIds.length);
                for (long nodeId : nodeIds) {
                    if (row.get(nodeId) == null) {
                        // If node not exist, just skip
                        continue;
                    }
                    IntFloatVector feat = ((GraphNode) (row.get(nodeId))).getFeats();
                    if (feat != null) {
                        nodeIdToFeats.put(nodeId, feat);
                    }
                }
                return new PartGetNodeFeatsResult(param.getPartKey().getPartitionId(), nodeIdToFeats);
            }
        default:
            {
                // TODO: support String, Int, and Any type node id
                throw new InvalidParameterException("Unsupport index type " + keyPart.getKeyType());
            }
    }
}

Example 25 with KeyPart

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

the class ReadTag method partitionGet.

@Override
public PartitionGetResult partitionGet(PartitionGetParam partParam) {
    GeneralPartGetParam param = (GeneralPartGetParam) partParam;
    KeyPart keyPart = param.getIndicesPart();
    long[] nodeIds = ((ILongKeyPartOp) keyPart).getKeys();
    ServerLongIntRow row = GraphMatrixUtils.getPSLongKeyIntRow(psContext, param, param.getRowId());
    DynamicLongArray nodes = new DynamicLongArray(nodeIds.length);
    for (int i = 0; i < nodeIds.length; i++) {
        long nodeId = nodeIds[i];
        if (row.get(nodeId) > 0) {
            nodes.add(nodeId);
        }
    }
    return new PartReadTagResult(nodes.getData());
}