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

Example 1 with KeyPart

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

the class UserRequestAdapter method get.

private FutureResult<Vector[]> get(IndexGetRowsRequest request) {
    // Only support column-partitioned matrix now!!
    MatrixMeta matrixMeta = PSAgentContext.get().getMatrixMetaManager().getMatrixMeta(request.getMatrixId());
    PartitionKey[] matrixParts = matrixMeta.getPartitionKeys();
    for (PartitionKey part : matrixParts) {
        if (part.getStartRow() != 0 || part.getEndRow() != matrixMeta.getRowNum()) {
            throw new UnsupportedOperationException("Get rows by indices only support column-partitioned matrix now");
        }
    }
    // Split the user request to partition requests
    FutureResult<Vector[]> result = new FutureResult<>();
    long startTs = System.currentTimeMillis();
    KeyPart[] splits;
    if (request instanceof IntIndexGetRowsRequest) {
        splits = RouterUtils.split(matrixMeta, -1, ((IntIndexGetRowsRequest) request).getIndices());
    } else if (request instanceof LongIndexGetRowsRequest) {
        splits = RouterUtils.split(matrixMeta, -1, ((LongIndexGetRowsRequest) request).getIndices());
    } else {
        throw new UnsupportedOperationException("Unsupport index request type " + request.getClass().toString());
    }
    LOG.info("Get by indices split use time: " + (System.currentTimeMillis() - startTs));
    assert matrixParts.length == splits.length;
    // filter empty partition requests
    int needRequestPartNum = noEmptyPartNum(splits);
    if (needRequestPartNum == 0) {
        result.set(null);
        return result;
    }
    // Create partition results cache
    ResponseCache cache = new MapResponseCache(needRequestPartNum);
    int requestId = request.getRequestId();
    requestIdToResponseCache.put(requestId, cache);
    requestIdToResultMap.put(requestId, result);
    requests.put(requestId, request);
    MatrixTransportClient matrixClient = PSAgentContext.get().getMatrixTransportClient();
    for (int i = 0; i < splits.length; i++) {
        if (splits[i] != null && splits[i].size() > 0) {
            sendIndexGetRowsRequest(matrixClient, requestId, request.getMatrixId(), request.getRowIds(), matrixParts[i].getPartitionId(), splits[i], request.getFunc());
        }
    }
    return result;
}

Example 2 with KeyPart

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

the class GetNodeTypes 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<int[]> nodeIdToTypes = new Long2ObjectOpenHashMap<>(nodeIds.length);
                for (long nodeId : nodeIds) {
                    if (row.get(nodeId) == null) {
                        // If node not exist, just skip
                        continue;
                    }
                    int[] types = ((GraphNode) (row.get(nodeId))).getTypes();
                    if (types != null) {
                        nodeIdToTypes.put(nodeId, types);
                    }
                }
                return new PartGetIntArrayAttrsResult(param.getPartKey().getPartitionId(), nodeIdToTypes);
            }
        default:
            {
                // TODO: support String, Int, and Any type node id
                throw new InvalidParameterException("Unsupport index type " + keyPart.getKeyType());
            }
    }
}

Example 3 with KeyPart

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

the class GetEdgeFeats 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[] edgeFeats = ((GraphNode) (row.get(nodeId))).getEdgeFeatures();
                    if (edgeFeats != null) {
                        nodeIdToFeats.put(nodeId, edgeFeats);
                    }
                }
                return new PartGetEdgeFeatsResult(param.getPartKey().getPartitionId(), nodeIdToFeats);
            }
        default:
            {
                // TODO: support String, Int, and Any type node id
                throw new InvalidParameterException("Unsupport index type " + keyPart.getKeyType());
            }
    }
}

Example 4 with KeyPart

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

the class GetEdgeTypes 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<int[]> nodeIdToTypes = new Long2ObjectOpenHashMap<>(nodeIds.length);
                for (long nodeId : nodeIds) {
                    if (row.get(nodeId) == null) {
                        // If node not exist, just skip
                        continue;
                    }
                    int[] types = ((GraphNode) (row.get(nodeId))).getEdgeTypes();
                    if (types != null) {
                        nodeIdToTypes.put(nodeId, types);
                    }
                }
                return new PartGetIntArrayAttrsResult(param.getPartKey().getPartitionId(), nodeIdToTypes);
            }
        default:
            {
                // TODO: support String, Int, and Any type node id
                throw new InvalidParameterException("Unsupport index type " + keyPart.getKeyType());
            }
    }
}

Example 5 with KeyPart

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

the class GetEdgeWeights 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[]> nodeIdToWeights = new Long2ObjectOpenHashMap<>(nodeIds.length);
                for (long nodeId : nodeIds) {
                    if (row.get(nodeId) == null) {
                        // If node not exist, just skip
                        continue;
                    }
                    float[] weights = ((GraphNode) (row.get(nodeId))).getWeights();
                    if (weights != null) {
                        nodeIdToWeights.put(nodeId, weights);
                    }
                }
                return new PartGetFloatArrayAttrsResult(param.getPartKey().getPartitionId(), nodeIdToWeights);
            }
        default:
            {
                // TODO: support String, Int, and Any type node id
                throw new InvalidParameterException("Unsupport index type " + keyPart.getKeyType());
            }
    }
}