Examples with OutputContext org.apache.tez.runtime.api.OutputContext used on opensource projects

Example 46 with OutputContext

use of org.apache.tez.runtime.api.OutputContext in project tez by apache.

the class OutputTestHelpers method createOutputContext.

static OutputContext createOutputContext(Configuration conf, Path workingDir) throws IOException {
    OutputContext ctx = mock(OutputContext.class);
    doAnswer(new Answer<Void>() {

        @Override
        public Void answer(InvocationOnMock invocation) throws Throwable {
            long requestedSize = (Long) invocation.getArguments()[0];
            MemoryUpdateCallbackHandler callback = (MemoryUpdateCallbackHandler) invocation.getArguments()[1];
            callback.memoryAssigned(requestedSize);
            return null;
        }
    }).when(ctx).requestInitialMemory(anyLong(), any(MemoryUpdateCallback.class));
    doReturn(TezUtils.createUserPayloadFromConf(conf)).when(ctx).getUserPayload();
    doReturn("destinationVertex").when(ctx).getDestinationVertexName();
    doReturn("UUID").when(ctx).getUniqueIdentifier();
    doReturn(new String[] { workingDir.toString() }).when(ctx).getWorkDirs();
    doReturn(200 * 1024 * 1024l).when(ctx).getTotalMemoryAvailableToTask();
    doReturn(new TezCounters()).when(ctx).getCounters();
    OutputStatisticsReporter statsReporter = mock(OutputStatisticsReporter.class);
    doReturn(statsReporter).when(ctx).getStatisticsReporter();
    doReturn(new ExecutionContextImpl("localhost")).when(ctx).getExecutionContext();
    return ctx;
}

Example 47 with OutputContext

use of org.apache.tez.runtime.api.OutputContext in project tez by apache.

the class OutputTestHelpers method createOutputContext.

static OutputContext createOutputContext() throws IOException {
    OutputContext outputContext = mock(OutputContext.class);
    Configuration conf = new TezConfiguration();
    UserPayload payLoad = TezUtils.createUserPayloadFromConf(conf);
    String[] workingDirs = new String[] { "workDir1" };
    OutputStatisticsReporter statsReporter = mock(OutputStatisticsReporter.class);
    TezCounters counters = new TezCounters();
    doReturn("destinationVertex").when(outputContext).getDestinationVertexName();
    doReturn(payLoad).when(outputContext).getUserPayload();
    doReturn(workingDirs).when(outputContext).getWorkDirs();
    doReturn(200 * 1024 * 1024l).when(outputContext).getTotalMemoryAvailableToTask();
    doReturn(counters).when(outputContext).getCounters();
    doReturn(statsReporter).when(outputContext).getStatisticsReporter();
    return outputContext;
}

Example 48 with OutputContext

use of org.apache.tez.runtime.api.OutputContext in project tez by apache.

the class TestOnFileUnorderedKVOutput method testGeneratedDataMovementEvent.

@Test(timeout = 5000)
public void testGeneratedDataMovementEvent() throws Exception {
    Configuration conf = new Configuration();
    conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_KEY_CLASS, Text.class.getName());
    conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_VALUE_CLASS, IntWritable.class.getName());
    TezSharedExecutor sharedExecutor = new TezSharedExecutor(conf);
    OutputContext outputContext = createOutputContext(conf, sharedExecutor);
    UnorderedKVOutput kvOutput = new UnorderedKVOutput(outputContext, 1);
    List<Event> events = null;
    events = kvOutput.initialize();
    kvOutput.start();
    assertTrue(events != null && events.size() == 0);
    KeyValueWriter kvWriter = kvOutput.getWriter();
    List<KVPair> data = KVDataGen.generateTestData(true, 0);
    for (KVPair kvp : data) {
        kvWriter.write(kvp.getKey(), kvp.getvalue());
    }
    events = kvOutput.close();
    assertEquals(45, task.getTaskStatistics().getIOStatistics().values().iterator().next().getDataSize());
    assertEquals(5, task.getTaskStatistics().getIOStatistics().values().iterator().next().getItemsProcessed());
    assertTrue(events != null && events.size() == 2);
    CompositeDataMovementEvent dmEvent = (CompositeDataMovementEvent) events.get(1);
    assertEquals("Invalid source index", 0, dmEvent.getSourceIndexStart());
    DataMovementEventPayloadProto shufflePayload = DataMovementEventPayloadProto.parseFrom(ByteString.copyFrom(dmEvent.getUserPayload()));
    assertFalse(shufflePayload.hasEmptyPartitions());
    assertEquals(outputContext.getUniqueIdentifier(), shufflePayload.getPathComponent());
    assertEquals(shufflePort, shufflePayload.getPort());
    assertEquals("localhost", shufflePayload.getHost());
    sharedExecutor.shutdownNow();
}

Example 49 with OutputContext

use of org.apache.tez.runtime.api.OutputContext in project tez by apache.

the class TestUnorderedKVOutput2 method testClose.

@Test(timeout = 10000)
public void testClose() throws Exception {
    OutputContext outputContext = OutputTestHelpers.createOutputContext(conf, workingDir);
    int numPartitions = 1;
    UnorderedKVOutput output = new UnorderedKVOutput(outputContext, numPartitions);
    output.initialize();
    output.start();
    assertNotNull(output.getWriter());
    output.close();
    assertNull(output.getWriter());
}

Example 50 with OutputContext

use of org.apache.tez.runtime.api.OutputContext in project tez by apache.

the class TestUnorderedKVOutput2 method testNonStartedOutput.

@Test(timeout = 5000)
public void testNonStartedOutput() throws Exception {
    OutputContext outputContext = OutputTestHelpers.createOutputContext();
    int numPartitions = 1;
    UnorderedKVOutput output = new UnorderedKVOutput(outputContext, numPartitions);
    output.initialize();
    List<Event> events = output.close();
    assertEquals(1, events.size());
    Event event1 = events.get(0);
    assertTrue(event1 instanceof DataMovementEvent);
    DataMovementEvent dme = (DataMovementEvent) event1;
    ByteBuffer bb = dme.getUserPayload();
    ShuffleUserPayloads.DataMovementEventPayloadProto shufflePayload = ShuffleUserPayloads.DataMovementEventPayloadProto.parseFrom(ByteString.copyFrom(bb));
    assertTrue(shufflePayload.hasEmptyPartitions());
    byte[] emptyPartitions = TezCommonUtils.decompressByteStringToByteArray(shufflePayload.getEmptyPartitions());
    BitSet emptyPartionsBitSet = TezUtilsInternal.fromByteArray(emptyPartitions);
    assertEquals(numPartitions, emptyPartionsBitSet.cardinality());
    for (int i = 0; i < numPartitions; i++) {
        assertTrue(emptyPartionsBitSet.get(i));
    }
}