@Override
  public GroupingShuffleReaderIterator<K, V> iterator() throws IOException {
    ApplianceShuffleEntryReader entryReader =
        new ApplianceShuffleEntryReader(
            shuffleReaderConfig, executionContext, operationContext, true);

    initCounter(entryReader.getDatasetId());

    return iterator(entryReader);
  }
@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(
          newShuffleEntry(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().toByteArray());
    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 /* do not sort values */);

    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))));
      int numRecordsReturned = 1; // including start() above.
      for (; iter.advance(); ++numRecordsReturned) {
        iter.getCurrent().getValue(); // ignored
      }
      assertEquals(kNumRecords, numRecordsReturned);

      // Cannot split since all input was consumed.
      assertNull(
          iter.requestDynamicSplit(splitRequestAtPosition(makeShufflePosition(kFirstShard, null))));
    }
    assertTrue(shuffleReader.isClosed());
  }