use of com.tencent.angel.ml.math2.vector.IntKeyVector in project angel by Tencent.
the class UpdateColsParam method split.
@Override
public List<PartitionUpdateParam> split() {
List<PartitionKey> pkeys = PSAgentContext.get().getMatrixMetaManager().getPartitions(matrixId);
List<PartitionUpdateParam> params = new ArrayList<>();
int start = 0, end = 0;
for (PartitionKey pkey : pkeys) {
long startCol = pkey.getStartCol();
long endCol = pkey.getEndCol();
if (start < ((IntKeyVector) cols).getDim() && VectorUtils.getLong(cols, start) >= startCol) {
while (end < ((IntKeyVector) cols).getDim() && VectorUtils.getLong(cols, end) < endCol) end++;
long[] part = new long[end - start];
if (cols instanceof IntIntVector) {
ArrayCopy.copy(((IntIntVector) cols).getStorage().getValues(), start, part, 0, end - start);
} else {
System.arraycopy(((IntLongVector) cols).getStorage().getValues(), start, part, 0, end - start);
}
long firstKey = 0l;
for (Map.Entry<Long, Vector> first : values.entrySet()) {
firstKey = first.getKey();
break;
}
if (values.get(firstKey) instanceof IntDoubleVector) {
IntDoubleVector[] updates = new IntDoubleVector[part.length];
for (int i = 0; i < part.length; i++) updates[i] = (IntDoubleVector) values.get(part[i]);
params.add(new PartitionUpdateColsParam(matrixId, pkey, rows, part, VFactory.compIntDoubleVector(rows.length, updates, part.length), op));
} else if (values.get(firstKey) instanceof IntFloatVector) {
IntFloatVector[] updates = new IntFloatVector[part.length];
for (int i = 0; i < part.length; i++) updates[i] = (IntFloatVector) values.get(part[i]);
params.add(new PartitionUpdateColsParam(matrixId, pkey, rows, part, VFactory.compIntFloatVector(rows.length, updates, part.length), op));
} else {
throw new AngelException("Update data type should be float or double!");
}
start = end;
}
}
return params;
}
Aggregations