GroupingShuffleReader extends NativeReader<WindowedValue<KV<K, Reiterable<V>>>> { @Override public GroupingShuffleReaderIterator<K, V> iterator() throws IOException { ApplianceShuffleEntryReader entryReader = new ApplianceShuffleEntryReader( shuffleReaderConfig, executionContext, operationContext, true); initCounter(entryReader.getDatasetId()); return iterator(entryReader); } GroupingShuffleReader(
      PipelineOptions options,
      byte[] shuffleReaderConfig,
      @Nullable String startShufflePosition,
      @Nullable String stopShufflePosition,
      Coder<WindowedValue<KV<K, Iterable<V>>>> coder,
      BatchModeExecutionContext executionContext,
      DataflowOperationContext operationContext,
      ShuffleReadCounterFactory shuffleReadCounterFactory,
      boolean valuesAreSorted); @Override GroupingShuffleReaderIterator<K, V> iterator();  }
@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)); } 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)) { assertTrue(iter.start()); assertNull( iter.requestDynamicSplit( splitRequestAtPosition(makeShufflePosition(kNumRecords + 1, null)))); byte[] key = CoderUtils.encodeToByteArray(BigEndianIntegerCoder.of(), 0); assertNull( iter.requestDynamicSplit(splitRequestAtPosition(makeShufflePosition(kFirstShard, key)))); assertNull( iter.requestDynamicSplit(splitRequestAtPosition(makeShufflePosition(kFirstShard, null)))); int numRecordsReturned = 1; for (; iter.advance(); ++numRecordsReturned) { iter.getCurrent().getValue(); } assertEquals(kNumRecords, numRecordsReturned); assertNull( iter.requestDynamicSplit(splitRequestAtPosition(makeShufflePosition(kFirstShard, null)))); } assertTrue(shuffleReader.isClosed()); }
@Test public void testConsumedParallelism() throws Exception { PipelineOptions options = PipelineOptionsFactory.create(); BatchModeExecutionContext context = BatchModeExecutionContext.forTesting(options, "testStage"); final int kFirstShard = 0; TestShuffleReader shuffleReader = new TestShuffleReader(); final int kNumRecords = 5; for (int i = 0; i < kNumRecords; ++i) { byte[] key = CoderUtils.encodeToByteArray(BigEndianIntegerCoder.of(), i); ShuffleEntry entry = new ShuffleEntry(fabricatePosition(kFirstShard, i), key, EMPTY_BYTE_ARRAY, key); shuffleReader.addEntry(entry); } 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 ); assertFalse(shuffleReader.isClosed()); try (GroupingShuffleReaderIterator<Integer, Integer> iter = groupingShuffleReader.iterator(shuffleReader)) { assertEquals(0.0, consumedParallelismFromProgress(iter.getProgress()), 0); assertTrue(iter.start()); assertEquals( 0.0, readerProgressToCloudProgress(iter.getProgress()).getConsumedParallelism().getValue(), 0); assertNotNull( iter.requestDynamicSplit( splitRequestAtPosition( makeShufflePosition( fabricatePosition(kFirstShard, 2).immediateSuccessor().getPosition())))); assertEquals(0.0, consumedParallelismFromProgress(iter.getProgress()), 0); assertTrue(iter.advance()); assertEquals(1.0, consumedParallelismFromProgress(iter.getProgress()), 0); assertTrue(iter.advance()); assertEquals(2.0, consumedParallelismFromProgress(iter.getProgress()), 0); assertFalse(iter.advance()); assertEquals(3.0, consumedParallelismFromProgress(iter.getProgress()), 0); } assertTrue(shuffleReader.isClosed()); }
@Test public void testReadFromShuffleAndDynamicSplit() throws Exception { 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 ); groupingShuffleReader.perOperationPerDatasetBytesCounter = operationContext .counterFactory() .longSum(CounterName.named("dax-shuffle-test-wf-read-bytes")); TestShuffleReader shuffleReader = new TestShuffleReader(); final int kNumRecords = 10; final int kFirstShard = 0; final int kSecondShard = 1; for (int i = 0; i < kNumRecords; ++i) { byte[] keyByte = CoderUtils.encodeToByteArray(BigEndianIntegerCoder.of(), i); ShuffleEntry entry = new ShuffleEntry( fabricatePosition(kFirstShard, keyByte), keyByte, EMPTY_BYTE_ARRAY, keyByte); shuffleReader.addEntry(entry); } for (int i = kNumRecords; i < 2 * kNumRecords; ++i) { byte[] keyByte = CoderUtils.encodeToByteArray(BigEndianIntegerCoder.of(), i); ShuffleEntry entry = new ShuffleEntry( fabricatePosition(kSecondShard, keyByte), keyByte, EMPTY_BYTE_ARRAY, keyByte); shuffleReader.addEntry(entry); } int i = 0; assertFalse(shuffleReader.isClosed()); try (GroupingShuffleReaderIterator<Integer, Integer> iter = groupingShuffleReader.iterator(shuffleReader)) { assertTrue(iter.start()); ++i; assertNull(iter.requestDynamicSplit(splitRequestAtPosition(new Position()))); NativeReader.DynamicSplitResult dynamicSplitResult = iter.requestDynamicSplit(splitRequestAtPosition(makeShufflePosition(kSecondShard, null))); assertNotNull(dynamicSplitResult); assertEquals( encodeBase64URLSafeString(fabricatePosition(kSecondShard).getPosition()), positionFromSplitResult(dynamicSplitResult).getShufflePosition()); for (; iter.advance(); ++i) { iter.getCurrent(); iter.getCurrent(); KV<Integer, Reiterable<Integer>> elem = iter.getCurrent().getValue(); int key = elem.getKey(); assertEquals(key, i); Reiterable<Integer> valuesIterable = elem.getValue(); Reiterator<Integer> valuesIterator = valuesIterable.iterator(); int j = 0; while (valuesIterator.hasNext()) { assertTrue(valuesIterator.hasNext()); assertTrue(valuesIterator.hasNext()); int value = valuesIterator.next(); assertEquals(value, i); ++j; } assertFalse(valuesIterator.hasNext()); assertFalse(valuesIterator.hasNext()); assertEquals(1, j); } assertFalse(iter.advance()); } assertTrue(shuffleReader.isClosed()); assertEquals(i, kNumRecords); assertEquals( 80L, (long) groupingShuffleReader.perOperationPerDatasetBytesCounter.getAggregate()); }
@Test public void testGetApproximateProgress() throws Exception { 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()); 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()); Position proposedSplitPosition = new Position(); String stop = encodeBase64URLSafeString(fabricatePosition(0).getPosition()); proposedSplitPosition.setShufflePosition(stop); assertNull( iter.requestDynamicSplit( toDynamicSplitRequest(approximateSplitRequestAtPosition(proposedSplitPosition)))); } assertTrue(shuffleReader.isClosed()); }
@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(); 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)) { assertTrue(iter.start()); int numRecordsReturned = 1; for (; iter.advance(); ++numRecordsReturned) { if (numRecordsReturned > 5) { setCurrentExecutionState(MOCK_ORIGINAL_NAME_FOR_EXECUTING_STEP2); } 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); }