Examples with CounterBackedElementByteSizeObserver org.apache.beam.runners.dataflow.worker.counters.CounterBackedElementByteSizeObserver used on opensource projects

Example 1 with CounterBackedElementByteSizeObserver

use of org.apache.beam.runners.dataflow.worker.counters.CounterBackedElementByteSizeObserver in project beam by apache.

the class OutputObjectAndByteCounter method countMeanByte.

/**
 * Count output mean byte.
 */
public OutputObjectAndByteCounter countMeanByte(String meanByteCounterName) {
    if (elementByteSizeObservable != null) {
        meanByteCount = counterFactory.longMean(getCounterName(meanByteCounterName));
        meanByteCountObserver = new CounterBackedElementByteSizeObserver(meanByteCount);
    }
    return this;
}

Example 2 with CounterBackedElementByteSizeObserver

use of org.apache.beam.runners.dataflow.worker.counters.CounterBackedElementByteSizeObserver in project beam by apache.

the class GroupingShuffleReaderTest method runIterationOverGroupingShuffleReader.

@SuppressWarnings("ReturnValueIgnored")
private List<KV<Integer, List<KV<Integer, Integer>>>> runIterationOverGroupingShuffleReader(BatchModeExecutionContext context, TestShuffleReader shuffleReader, GroupingShuffleReader<Integer, KV<Integer, Integer>> groupingShuffleReader, Coder<WindowedValue<KV<Integer, Iterable<KV<Integer, Integer>>>>> coder, ValuesToRead valuesToRead) throws Exception {
    CounterSet counterSet = new CounterSet();
    Counter<Long, ?> elementByteSizeCounter = counterSet.longSum(CounterName.named("element-byte-size-counter"));
    CounterBackedElementByteSizeObserver elementObserver = new CounterBackedElementByteSizeObserver(elementByteSizeCounter);
    List<KV<Integer, List<KV<Integer, Integer>>>> actual = new ArrayList<>();
    assertFalse(shuffleReader.isClosed());
    try (GroupingShuffleReaderIterator<Integer, KV<Integer, Integer>> iter = groupingShuffleReader.iterator(shuffleReader)) {
        Iterable<KV<Integer, Integer>> prevValuesIterable = null;
        Iterator<KV<Integer, Integer>> prevValuesIterator = null;
        for (boolean more = iter.start(); more; more = iter.advance()) {
            // Should not fail.
            iter.getCurrent();
            iter.getCurrent();
            // safe co-variant cast from Reiterable to Iterable
            @SuppressWarnings({ // TODO(https://issues.apache.org/jira/browse/BEAM-10556)
            "rawtypes", "unchecked" }) WindowedValue<KV<Integer, Iterable<KV<Integer, Integer>>>> windowedValue = (WindowedValue) iter.getCurrent();
            // Verify that the byte size observer is lazy for every value the GroupingShuffleReader
            // produces.
            coder.registerByteSizeObserver(windowedValue, elementObserver);
            assertTrue(elementObserver.getIsLazy());
            // Verify value is in an empty windows.
            assertEquals(BoundedWindow.TIMESTAMP_MIN_VALUE, windowedValue.getTimestamp());
            assertEquals(0, windowedValue.getWindows().size());
            KV<Integer, Iterable<KV<Integer, Integer>>> elem = windowedValue.getValue();
            Integer key = elem.getKey();
            List<KV<Integer, Integer>> values = new ArrayList<>();
            if (valuesToRead.ordinal() > ValuesToRead.SKIP_VALUES.ordinal()) {
                if (prevValuesIterable != null) {
                    // Verifies that this does not throw.
                    prevValuesIterable.iterator();
                }
                if (prevValuesIterator != null) {
                    // Verifies that this does not throw.
                    prevValuesIterator.hasNext();
                }
                Iterable<KV<Integer, Integer>> valuesIterable = elem.getValue();
                Iterator<KV<Integer, Integer>> valuesIterator = valuesIterable.iterator();
                if (valuesToRead.ordinal() >= ValuesToRead.READ_ONE_VALUE.ordinal()) {
                    while (valuesIterator.hasNext()) {
                        assertTrue(valuesIterator.hasNext());
                        assertTrue(valuesIterator.hasNext());
                        assertEquals("BatchModeExecutionContext key", key, context.getKey());
                        values.add(valuesIterator.next());
                        if (valuesToRead == ValuesToRead.READ_ONE_VALUE) {
                            break;
                        }
                    }
                    if (valuesToRead.ordinal() >= ValuesToRead.READ_ALL_VALUES.ordinal()) {
                        assertFalse(valuesIterator.hasNext());
                        assertFalse(valuesIterator.hasNext());
                        try {
                            valuesIterator.next();
                            fail("Expected NoSuchElementException");
                        } catch (NoSuchElementException exn) {
                        // As expected.
                        }
                        // Verifies that this does not throw.
                        valuesIterable.iterator();
                    }
                }
                if (valuesToRead == ValuesToRead.READ_ALL_VALUES_TWICE) {
                    // Create new iterator;
                    valuesIterator = valuesIterable.iterator();
                    while (valuesIterator.hasNext()) {
                        assertTrue(valuesIterator.hasNext());
                        assertTrue(valuesIterator.hasNext());
                        assertEquals("BatchModeExecutionContext key", key, context.getKey());
                        valuesIterator.next();
                    }
                    assertFalse(valuesIterator.hasNext());
                    assertFalse(valuesIterator.hasNext());
                    try {
                        valuesIterator.next();
                        fail("Expected NoSuchElementException");
                    } catch (NoSuchElementException exn) {
                    // As expected.
                    }
                }
                prevValuesIterable = valuesIterable;
                prevValuesIterator = valuesIterator;
            }
            actual.add(KV.of(key, values));
        }
        assertFalse(iter.advance());
        assertFalse(iter.advance());
        try {
            iter.getCurrent();
            fail("Expected NoSuchElementException");
        } catch (NoSuchElementException exn) {
        // As expected.
        }
    }
    assertTrue(shuffleReader.isClosed());
    return actual;
}

Example 3 with CounterBackedElementByteSizeObserver

use of org.apache.beam.runners.dataflow.worker.counters.CounterBackedElementByteSizeObserver in project beam by apache.

the class OutputObjectAndByteCounter method countBytes.

/**
 * Count output bytes.
 */
public OutputObjectAndByteCounter countBytes(String bytesCounterName) {
    if (elementByteSizeObservable != null) {
        byteCount = counterFactory.longSum(getCounterName(bytesCounterName));
        byteCountObserver = new CounterBackedElementByteSizeObserver(byteCount);
    }
    return this;
}