TransformTranslator { @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(); } }; } }; } private  TransformTranslator();   }
@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(1))), CoderHelpers.toByteArray( WindowedValue.of(2, now.plus(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(2))), CoderHelpers.toByteArray( WindowedValue.of(4, now.plus(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) { assertEquals( Arrays.asList( WindowedValue.of(KV.of(1, 1), now, window, paneInfo), WindowedValue.of(KV.of(1, 2), now.plus(1), window, paneInfo)), list); } else { assertEquals( Arrays.asList( WindowedValue.of(KV.of(2, 3), now, window, paneInfo), WindowedValue.of(KV.of(2, 4), now.plus(2), window, paneInfo)), list); } } }