@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 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 /* do not sort values */);
    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;

    // Setting up two shards with kNumRecords each; keys are unique
    // (hence groups of values for the same key are singletons)
    // therefore each record comes with a unique position constructed.
    for (int i = 0; i < kNumRecords; ++i) {
      byte[] keyByte = CoderUtils.encodeToByteArray(BigEndianIntegerCoder.of(), i);
      ShuffleEntry entry =
          newShuffleEntry(
              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 =
          newShuffleEntry(
              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)) {
      // Poke the iterator so we can test dynamic splitting.
      assertTrue(iter.start());
      ++i;

      assertNull(iter.requestDynamicSplit(splitRequestAtPosition(new Position())));

      // Split at the shard boundary
      NativeReader.DynamicSplitResult dynamicSplitResult =
          iter.requestDynamicSplit(splitRequestAtPosition(makeShufflePosition(kSecondShard, null)));
      assertNotNull(dynamicSplitResult);
      assertEquals(
          encodeBase64URLSafeString(fabricatePosition(kSecondShard).getPosition().toByteArray()),
          positionFromSplitResult(dynamicSplitResult).getShufflePosition());

      for (; iter.advance(); ++i) {
        // iter.getCurrent() is supposed to be side-effect-free and give the same result if called
        // repeatedly. Test that this is indeed the case.
        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);
    // There are 10 Shuffle records that each encode an integer key (4 bytes) and integer value (4
    // bytes). We therefore expect to read 80 bytes.
    assertEquals(
        80L, (long) groupingShuffleReader.perOperationPerDatasetBytesCounter.getAggregate());
  }