Examples with ShuffleReadCounter org.apache.beam.runners.dataflow.worker.util.common.worker.ShuffleReadCounter used on opensource projects

Example 1 with ShuffleReadCounter

use of org.apache.beam.runners.dataflow.worker.util.common.worker.ShuffleReadCounter in project beam by apache.

the class GroupingShuffleReaderTest method expectShuffleReadCounterEquals.

private void expectShuffleReadCounterEquals(TestShuffleReadCounterFactory factory, Map<String, Long> expectedReadBytesForOriginal) {
    ShuffleReadCounter src = factory.getOnlyShuffleReadCounterOrNull();
    assertNotNull(src);
    // should not be set.
    if (src.legacyPerOperationPerDatasetBytesCounter != null) {
        assertEquals(0, (long) src.legacyPerOperationPerDatasetBytesCounter.getAggregate());
    }
    // Verify that each executing step used when reading from the GroupingShuffleReader
    // has a counter with a bytes read value.
    assertEquals(expectedReadBytesForOriginal.size(), (long) src.counterSet.size());
    Iterator it = expectedReadBytesForOriginal.entrySet().iterator();
    while (it.hasNext()) {
        Map.Entry<String, Long> pair = (Map.Entry) it.next();
        Counter counter = src.counterSet.getExistingCounter(ShuffleReadCounter.generateCounterName(ORIGINAL_SHUFFLE_STEP_NAME, pair.getKey()));
        assertEquals(pair.getValue(), counter.getAggregate());
    }
}

Example 2 with ShuffleReadCounter

use of org.apache.beam.runners.dataflow.worker.util.common.worker.ShuffleReadCounter in project beam by apache.

the class TestShuffleReadCounterFactory method create.

@Override
public ShuffleReadCounter create(String originalShuffleStepName, boolean experimentEnabled, Counter<Long, Long> legacyPerOperationPerDatasetBytesCounter) {
    ShuffleReadCounter src = new ShuffleReadCounter(originalShuffleStepName, experimentEnabled, legacyPerOperationPerDatasetBytesCounter);
    originalShuffleStepToCounter.put(originalShuffleStepName, src);
    return src;
}