Search in sources :

Example 1 with ByteArrayShufflePosition

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

the class GroupingShuffleReaderTest method testGetApproximateProgress.

@Test
public void testGetApproximateProgress() throws Exception {
    // Store the positions of all KVs returned.
    List<ByteArrayShufflePosition> positionsList = new ArrayList<>();
    PipelineOptions options = PipelineOptionsFactory.create();
    BatchModeExecutionContext context = BatchModeExecutionContext.forTesting(options, "testStage");
    TestOperationContext operationContext = TestOperationContext.create();
    GroupingShuffleReader<Integer, Integer> groupingShuffleReader = new GroupingShuffleReader<>(options, null, null, null, WindowedValue.getFullCoder(KvCoder.of(BigEndianIntegerCoder.of(), IterableCoder.of(BigEndianIntegerCoder.of())), IntervalWindow.getCoder()), context, operationContext, ShuffleReadCounterFactory.INSTANCE, false);
    TestShuffleReader shuffleReader = new TestShuffleReader();
    final int kNumRecords = 10;
    for (int i = 0; i < kNumRecords; ++i) {
        ByteArrayShufflePosition position = fabricatePosition(i);
        byte[] keyByte = CoderUtils.encodeToByteArray(BigEndianIntegerCoder.of(), i);
        positionsList.add(position);
        ShuffleEntry entry = new ShuffleEntry(position, keyByte, EMPTY_BYTE_ARRAY, keyByte);
        shuffleReader.addEntry(entry);
    }
    assertFalse(shuffleReader.isClosed());
    try (GroupingShuffleReaderIterator<Integer, Integer> iter = groupingShuffleReader.iterator(shuffleReader)) {
        Integer i = 0;
        for (boolean more = iter.start(); more; more = iter.advance()) {
            ApproximateReportedProgress progress = readerProgressToCloudProgress(iter.getProgress());
            assertNotNull(progress.getPosition().getShufflePosition());
            // Compare returned position with the expected position.
            assertEquals(positionsList.get(i).encodeBase64(), progress.getPosition().getShufflePosition());
            WindowedValue<KV<Integer, Reiterable<Integer>>> elem = iter.getCurrent();
            assertEquals(i, elem.getValue().getKey());
            i++;
        }
        assertFalse(iter.advance());
        // Cannot split since all input was consumed.
        Position proposedSplitPosition = new Position();
        String stop = encodeBase64URLSafeString(fabricatePosition(0).getPosition());
        proposedSplitPosition.setShufflePosition(stop);
        assertNull(iter.requestDynamicSplit(toDynamicSplitRequest(approximateSplitRequestAtPosition(proposedSplitPosition))));
    }
    assertTrue(shuffleReader.isClosed());
}
Also used : ByteArrayShufflePosition(org.apache.beam.runners.dataflow.worker.util.common.worker.ByteArrayShufflePosition) ReaderTestUtils.approximateSplitRequestAtPosition(org.apache.beam.runners.dataflow.worker.ReaderTestUtils.approximateSplitRequestAtPosition) ByteArrayShufflePosition(org.apache.beam.runners.dataflow.worker.util.common.worker.ByteArrayShufflePosition) Position(com.google.api.services.dataflow.model.Position) ReaderTestUtils.splitRequestAtPosition(org.apache.beam.runners.dataflow.worker.ReaderTestUtils.splitRequestAtPosition) ArrayList(java.util.ArrayList) KV(org.apache.beam.sdk.values.KV) Base64.encodeBase64URLSafeString(com.google.api.client.util.Base64.encodeBase64URLSafeString) ShuffleEntry(org.apache.beam.runners.dataflow.worker.util.common.worker.ShuffleEntry) PipelineOptions(org.apache.beam.sdk.options.PipelineOptions) ApproximateReportedProgress(com.google.api.services.dataflow.model.ApproximateReportedProgress) Test(org.junit.Test)

Aggregations

Base64.encodeBase64URLSafeString (com.google.api.client.util.Base64.encodeBase64URLSafeString)1 ApproximateReportedProgress (com.google.api.services.dataflow.model.ApproximateReportedProgress)1 Position (com.google.api.services.dataflow.model.Position)1 ArrayList (java.util.ArrayList)1 ReaderTestUtils.approximateSplitRequestAtPosition (org.apache.beam.runners.dataflow.worker.ReaderTestUtils.approximateSplitRequestAtPosition)1 ReaderTestUtils.splitRequestAtPosition (org.apache.beam.runners.dataflow.worker.ReaderTestUtils.splitRequestAtPosition)1 ByteArrayShufflePosition (org.apache.beam.runners.dataflow.worker.util.common.worker.ByteArrayShufflePosition)1 ShuffleEntry (org.apache.beam.runners.dataflow.worker.util.common.worker.ShuffleEntry)1 PipelineOptions (org.apache.beam.sdk.options.PipelineOptions)1 KV (org.apache.beam.sdk.values.KV)1 Test (org.junit.Test)1