Examples with Base64.encodeBase64URLSafeString com.google.api.client.util.Base64.encodeBase64URLSafeString used on opensource projects

Example 1 with Base64.encodeBase64URLSafeString

use of com.google.api.client.util.Base64.encodeBase64URLSafeString in project beam by apache.

the class GroupingShuffleReaderTest method testReadFromShuffleDataAndFailToSplit.

@Test
public void testReadFromShuffleDataAndFailToSplit() throws Exception {
    PipelineOptions options = PipelineOptionsFactory.create();
    BatchModeExecutionContext context = BatchModeExecutionContext.forTesting(options, "testStage");
    final int kFirstShard = 0;
    TestShuffleReader shuffleReader = new TestShuffleReader();
    final int kNumRecords = 2;
    for (int i = 0; i < kNumRecords; ++i) {
        byte[] key = CoderUtils.encodeToByteArray(BigEndianIntegerCoder.of(), i);
        shuffleReader.addEntry(new ShuffleEntry(fabricatePosition(kFirstShard, key), key, EMPTY_BYTE_ARRAY, key));
    }
    // Note that TestShuffleReader start/end positions are in the
    // space of keys not the positions (TODO: should probably always
    // use positions instead).
    String stop = encodeBase64URLSafeString(fabricatePosition(kNumRecords).getPosition());
    TestOperationContext operationContext = TestOperationContext.create();
    GroupingShuffleReader<Integer, Integer> groupingShuffleReader = new GroupingShuffleReader<>(options, null, null, stop, WindowedValue.getFullCoder(KvCoder.of(BigEndianIntegerCoder.of(), IterableCoder.of(BigEndianIntegerCoder.of())), IntervalWindow.getCoder()), context, operationContext, ShuffleReadCounterFactory.INSTANCE, false);
    assertFalse(shuffleReader.isClosed());
    try (GroupingShuffleReaderIterator<Integer, Integer> iter = groupingShuffleReader.iterator(shuffleReader)) {
        // Poke the iterator so we can test dynamic splitting.
        assertTrue(iter.start());
        // Cannot split since the value provided is past the current stop position.
        assertNull(iter.requestDynamicSplit(splitRequestAtPosition(makeShufflePosition(kNumRecords + 1, null))));
        byte[] key = CoderUtils.encodeToByteArray(BigEndianIntegerCoder.of(), 0);
        // Cannot split since the split position is identical with the position of the record
        // that was just returned.
        assertNull(iter.requestDynamicSplit(splitRequestAtPosition(makeShufflePosition(kFirstShard, key))));
        // Cannot split since the requested split position comes before current position
        assertNull(iter.requestDynamicSplit(splitRequestAtPosition(makeShufflePosition(kFirstShard, null))));
        // including start() above.
        int numRecordsReturned = 1;
        for (; iter.advance(); ++numRecordsReturned) {
            // ignored
            iter.getCurrent().getValue();
        }
        assertEquals(kNumRecords, numRecordsReturned);
        // Cannot split since all input was consumed.
        assertNull(iter.requestDynamicSplit(splitRequestAtPosition(makeShufflePosition(kFirstShard, null))));
    }
    assertTrue(shuffleReader.isClosed());
}

Example 2 with Base64.encodeBase64URLSafeString

use of com.google.api.client.util.Base64.encodeBase64URLSafeString 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());
}

Example 3 with Base64.encodeBase64URLSafeString

use of com.google.api.client.util.Base64.encodeBase64URLSafeString in project beam by apache.

the class GroupingShuffleReaderTest method testShuffleReadCounterMultipleExecutingSteps.

@Test
public void testShuffleReadCounterMultipleExecutingSteps() throws Exception {
    PipelineOptions options = PipelineOptionsFactory.create();
    options.as(DataflowPipelineDebugOptions.class).setExperiments(Lists.newArrayList(Experiment.IntertransformIO.getName()));
    BatchModeExecutionContext context = BatchModeExecutionContext.forTesting(options, "testStage");
    final int kFirstShard = 0;
    TestShuffleReader shuffleReader = new TestShuffleReader();
    final int kNumRecords = 10;
    for (int i = 0; i < kNumRecords; ++i) {
        byte[] key = CoderUtils.encodeToByteArray(BigEndianIntegerCoder.of(), i);
        shuffleReader.addEntry(new ShuffleEntry(fabricatePosition(kFirstShard, key), key, EMPTY_BYTE_ARRAY, key));
    }
    TestShuffleReadCounterFactory shuffleReadCounterFactory = new TestShuffleReadCounterFactory();
    // Note that TestShuffleReader start/end positions are in the
    // space of keys not the positions (TODO: should probably always
    // use positions instead).
    String stop = encodeBase64URLSafeString(fabricatePosition(kNumRecords).getPosition());
    TestOperationContext operationContext = TestOperationContext.create();
    GroupingShuffleReader<Integer, Integer> groupingShuffleReader = new GroupingShuffleReader<>(options, null, null, stop, WindowedValue.getFullCoder(KvCoder.of(BigEndianIntegerCoder.of(), IterableCoder.of(BigEndianIntegerCoder.of())), IntervalWindow.getCoder()), context, operationContext, shuffleReadCounterFactory, false);
    assertFalse(shuffleReader.isClosed());
    try (GroupingShuffleReaderIterator<Integer, Integer> iter = groupingShuffleReader.iterator(shuffleReader)) {
        // Poke the iterator so we can test dynamic splitting.
        assertTrue(iter.start());
        // including start() above.
        int numRecordsReturned = 1;
        for (; iter.advance(); ++numRecordsReturned) {
            if (numRecordsReturned > 5) {
                setCurrentExecutionState(MOCK_ORIGINAL_NAME_FOR_EXECUTING_STEP2);
            }
            // ignored
            iter.getCurrent().getValue();
        }
        assertEquals(kNumRecords, numRecordsReturned);
    }
    assertTrue(shuffleReader.isClosed());
    Map<String, Long> expectedReadBytesMap = new HashMap<>();
    expectedReadBytesMap.put(MOCK_ORIGINAL_NAME_FOR_EXECUTING_STEP1, 48L);
    expectedReadBytesMap.put(MOCK_ORIGINAL_NAME_FOR_EXECUTING_STEP2, 32L);
    expectShuffleReadCounterEquals(shuffleReadCounterFactory, expectedReadBytesMap);
}