@VisibleForTesting
  static <K, V> Iterator<Iterator<WindowedValue<KV<K, V>>>> splitBySameKey(
      Iterator<Tuple2<ByteArray, byte[]>> in, Coder<K> keyCoder, Coder<WindowedValue<V>> wvCoder) {

    return new AbstractIterator<Iterator<WindowedValue<KV<K, V>>>>() {

      @Nullable Tuple2<ByteArray, byte[]> read = null;

      @Override
      protected Iterator<WindowedValue<KV<K, V>>> computeNext() {
        readNext();
        if (read != null) {
          byte[] value = read._1().getValue();
          byte[] keyPart = Arrays.copyOfRange(value, 0, value.length - 8);
          K key = CoderHelpers.fromByteArray(keyPart, keyCoder);
          return createIteratorForKey(keyPart, key);
        }
        return endOfData();
      }

      private void readNext() {
        if (read == null) {
          if (in.hasNext()) {
            read = in.next();
          }
        }
      }

      private void consumed() {
        read = null;
      }

      private Iterator<WindowedValue<KV<K, V>>> createIteratorForKey(byte[] keyPart, K key) {

        return new AbstractIterator<WindowedValue<KV<K, V>>>() {
          @Override
          protected WindowedValue<KV<K, V>> computeNext() {
            readNext();
            if (read != null) {
              byte[] value = read._1().getValue();
              byte[] prefix = Arrays.copyOfRange(value, 0, value.length - 8);
              if (Arrays.equals(prefix, keyPart)) {
                WindowedValue<V> wv = CoderHelpers.fromByteArray(read._2(), wvCoder);
                consumed();
                return WindowedValue.of(
                    KV.of(key, wv.getValue()), wv.getTimestamp(), wv.getWindows(), wv.getPane());
              }
            }
            return endOfData();
          }
        };
      }
    };
  }
@Test
  public void testSplitBySameKey() {
    VarIntCoder coder = VarIntCoder.of();
    WindowedValue.WindowedValueCoder<Integer> wvCoder =
        WindowedValue.FullWindowedValueCoder.of(coder, GlobalWindow.Coder.INSTANCE);
    Instant now = Instant.now();
    List<GlobalWindow> window = Arrays.asList(GlobalWindow.INSTANCE);
    PaneInfo paneInfo = PaneInfo.NO_FIRING;
    List<Tuple2<ByteArray, byte[]>> firstKey =
        Arrays.asList(
            new Tuple2(
                new ByteArray(CoderHelpers.toByteArrayWithTs(1, coder, now)),
                CoderHelpers.toByteArray(WindowedValue.of(1, now, window, paneInfo), wvCoder)),
            new Tuple2(
                new ByteArray(
                    CoderHelpers.toByteArrayWithTs(1, coder, now.plus(Duration.millis(1)))),
                CoderHelpers.toByteArray(
                    WindowedValue.of(2, now.plus(Duration.millis(1)), window, paneInfo), wvCoder)));

    List<Tuple2<ByteArray, byte[]>> secondKey =
        Arrays.asList(
            new Tuple2(
                new ByteArray(CoderHelpers.toByteArrayWithTs(2, coder, now)),
                CoderHelpers.toByteArray(WindowedValue.of(3, now, window, paneInfo), wvCoder)),
            new Tuple2(
                new ByteArray(
                    CoderHelpers.toByteArrayWithTs(2, coder, now.plus(Duration.millis(2)))),
                CoderHelpers.toByteArray(
                    WindowedValue.of(4, now.plus(Duration.millis(2)), window, paneInfo), wvCoder)));

    Iterable<Tuple2<ByteArray, byte[]>> concat = Iterables.concat(firstKey, secondKey);
    Iterator<Iterator<WindowedValue<KV<Integer, Integer>>>> keySplit;
    keySplit = TransformTranslator.splitBySameKey(concat.iterator(), coder, wvCoder);

    for (int i = 0; i < 2; i++) {
      Iterator<WindowedValue<KV<Integer, Integer>>> iter = keySplit.next();
      List<WindowedValue<KV<Integer, Integer>>> list = new ArrayList<>();
      Iterators.addAll(list, iter);
      if (i == 0) {
        // first key
        assertEquals(
            Arrays.asList(
                WindowedValue.of(KV.of(1, 1), now, window, paneInfo),
                WindowedValue.of(KV.of(1, 2), now.plus(Duration.millis(1)), window, paneInfo)),
            list);
      } else {
        // second key
        assertEquals(
            Arrays.asList(
                WindowedValue.of(KV.of(2, 3), now, window, paneInfo),
                WindowedValue.of(KV.of(2, 4), now.plus(Duration.millis(2)), window, paneInfo)),
            list);
      }
    }
  }