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());
}
Aggregations