コード例 #1
0
 public UserPayload toUserPayload() {
   return UserPayload.create(
       ByteBuffer.wrap(
           ShuffleEdgeManagerConfigPayloadProto.newBuilder()
               .setNumSourceTaskOutputs(numSourceTaskOutputs)
               .setNumDestinationTasks(numDestinationTasks)
               .setBasePartitionRange(basePartitionRange)
               .setRemainderRangeForLastShuffler(remainderRangeForLastShuffler)
               .build()
               .toByteArray()));
 }
コード例 #2
0
ファイル: TestValuesIterator.java プロジェクト: Umreth/tez
 private InputContext createTezInputContext() {
   TezCounters counters = new TezCounters();
   InputContext inputContext = mock(InputContext.class);
   doReturn(1024 * 1024 * 100l).when(inputContext).getTotalMemoryAvailableToTask();
   doReturn(counters).when(inputContext).getCounters();
   doReturn(1).when(inputContext).getInputIndex();
   doReturn("srcVertex").when(inputContext).getSourceVertexName();
   doReturn(1).when(inputContext).getTaskVertexIndex();
   doReturn(UserPayload.create(ByteBuffer.wrap(new byte[1024])))
       .when(inputContext)
       .getUserPayload();
   return inputContext;
 }
コード例 #3
0
  @Test(timeout = 5000)
  public void testGetBytePayload() throws IOException {
    int numBuckets = 10;
    VertexManagerPluginContext context = mock(VertexManagerPluginContext.class);
    CustomVertexConfiguration vertexConf =
        new CustomVertexConfiguration(numBuckets, TezWork.VertexType.INITIALIZED_EDGES);
    DataOutputBuffer dob = new DataOutputBuffer();
    vertexConf.write(dob);
    UserPayload payload = UserPayload.create(ByteBuffer.wrap(dob.getData()));
    when(context.getUserPayload()).thenReturn(payload);

    CustomPartitionVertex vm = new CustomPartitionVertex(context);
    vm.initialize();

    // prepare empty routing table
    Multimap<Integer, Integer> routingTable = HashMultimap.<Integer, Integer>create();
    payload = vm.getBytePayload(routingTable);
    // get conf from user payload
    CustomEdgeConfiguration edgeConf = new CustomEdgeConfiguration();
    DataInputByteBuffer dibb = new DataInputByteBuffer();
    dibb.reset(payload.getPayload());
    edgeConf.readFields(dibb);
    assertEquals(numBuckets, edgeConf.getNumBuckets());
  }