/** * Add data to the cache. * * @param workerInfo the remote worker destination * @param partitionId the remote Partition this message belongs to * @param destVertexId vertex id that is ultimate destination * @param data Data to send to remote worker * @return Size of messages for the worker. */ public int addData(WorkerInfo workerInfo, int partitionId, I destVertexId, T data) { // Get the data collection VertexIdData<I, T> partitionData = getPartitionData(workerInfo, partitionId); int originalSize = partitionData.getSize(); partitionData.add(destVertexId, data); // Update the size of cached, outgoing data per worker return incrDataSize(workerInfo.getTaskId(), partitionData.getSize() - originalSize); }
@Override public void readFieldsRequest(DataInput input) throws IOException { int numPartitions = input.readInt(); partitionVertexData = new PairList<Integer, B>(); partitionVertexData.initialize(numPartitions); while (numPartitions-- > 0) { final int partitionId = input.readInt(); B vertexIdData = createVertexIdData(); vertexIdData.setConf(getConf()); vertexIdData.readFields(input); partitionVertexData.add(partitionId, vertexIdData); } }
/** * This method tries to get a partition data from the data cache. * If null, it will create one. * * @param workerInfo The remote worker destination * @param partitionId The remote Partition this message belongs to * @return The partition data in data cache */ private VertexIdData<I, T> getPartitionData(WorkerInfo workerInfo, int partitionId) { // Get the data collection B partitionData = getData(partitionId); if (partitionData == null) { partitionData = createVertexIdData(); partitionData.setConf(getConf()); partitionData.initialize(getInitialBufferSize(workerInfo.getTaskId())); setData(partitionId, partitionData); } return partitionData; } }
@Override public void writeRequest(DataOutput output) throws IOException { output.writeInt(partitionVertexData.getSize()); PairList<Integer, B>.Iterator iterator = partitionVertexData.getIterator(); while (iterator.hasNext()) { iterator.next(); output.writeInt(iterator.getCurrentFirst()); iterator.getCurrentSecond().write(output); } }
@Override public int getSerializedSize() { int size = super.getSerializedSize() + 4; PairList<Integer, B>.Iterator iterator = partitionVertexData.getIterator(); while (iterator.hasNext()) { iterator.next(); size += 4 + iterator.getCurrentSecond().getSerializedSize(); } return size; } }
/** * This method is similar to the method above, * but use a serialized id to replace original I type * destVertexId. * * @param workerInfo The remote worker destination * @param partitionId The remote Partition this message belongs to * @param serializedId The byte array to store the serialized target vertex id * @param idPos The length of bytes of serialized id in the byte array above * @param data Data to send to remote worker * @return The number of bytes added to the target worker */ public int addData(WorkerInfo workerInfo, int partitionId, byte[] serializedId, int idPos, T data) { // Get the data collection VertexIdData<I, T> partitionData = getPartitionData(workerInfo, partitionId); int originalSize = partitionData.getSize(); partitionData.add(serializedId, idPos, data); // Update the size of cached, outgoing data per worker return incrDataSize(workerInfo.getTaskId(), partitionData.getSize() - originalSize); }