KeyedPValueTrackingVisitor extends PipelineVisitor.Defaults { public Set<PValue> getKeyedPValues() { checkState( finalized, "can't call getKeyedPValues before a Pipeline has been completely traversed"); return keyedValues; } private  KeyedPValueTrackingVisitor(); static KeyedPValueTrackingVisitor create(); @Override CompositeBehavior enterCompositeTransform(TransformHierarchy.Node node); @Override void leaveCompositeTransform(TransformHierarchy.Node node); @Override void visitValue(PValue value, TransformHierarchy.Node producer); Set<PValue> getKeyedPValues();  }
@Test public void groupByKeyProducesKeyedOutput() { PCollection<KV<String, Iterable<Integer>>> keyed = p.apply(Create.of(KV.of("foo", 3))) .apply(new DirectGroupByKeyOnly<>()) .apply( new DirectGroupAlsoByWindow<>( WindowingStrategy.globalDefault(), WindowingStrategy.globalDefault())); p.traverseTopologically(visitor); assertThat(visitor.getKeyedPValues(), hasItem(keyed)); }
@Test public void noInputUnkeyedOutput() { PCollection<KV<Integer, Iterable<Void>>> unkeyed = p.apply( Create.of(KV.<Integer, Iterable<Void>>of(-1, Collections.emptyList())) .withCoder(KvCoder.of(VarIntCoder.of(), IterableCoder.of(VoidCoder.of())))); p.traverseTopologically(visitor); assertThat(visitor.getKeyedPValues(), not(hasItem(unkeyed))); }
@Test public void unkeyedInputWithKeyPreserving() { PCollection<KV<String, Iterable<WindowedValue<KV<String, Integer>>>>> input = p.apply( Create.of( KV.of( "hello", (Iterable<WindowedValue<KV<String, Integer>>>) Collections.<WindowedValue<KV<String, Integer>>>emptyList())) .withCoder( KvCoder.of( StringUtf8Coder.of(), IterableCoder.of( WindowedValue.getValueOnlyCoder( KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of())))))); PCollection<KeyedWorkItem<String, KV<String, Integer>>> unkeyed = input .apply(ParDo.of(new ParDoMultiOverrideFactory.ToKeyedWorkItem<>())) .setCoder( KeyedWorkItemCoder.of( StringUtf8Coder.of(), KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()), GlobalWindow.Coder.INSTANCE)); p.traverseTopologically(visitor); assertThat(visitor.getKeyedPValues(), not(hasItem(unkeyed))); }
@Test public void keyedInputWithKeyPreserving() { PCollection<KV<String, WindowedValue<KV<String, Integer>>>> input = p.apply( Create.of( KV.of( "hello", WindowedValue.of( KV.of("hello", 3), new Instant(0), new IntervalWindow(new Instant(0), new Instant(9)), PaneInfo.NO_FIRING))) .withCoder( KvCoder.of( StringUtf8Coder.of(), WindowedValue.getValueOnlyCoder( KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()))))); TupleTag<KeyedWorkItem<String, KV<String, Integer>>> keyedTag = new TupleTag<>(); PCollection<KeyedWorkItem<String, KV<String, Integer>>> keyed = input .apply(new DirectGroupByKeyOnly<>()) .apply( new DirectGroupAlsoByWindow<>( WindowingStrategy.globalDefault(), WindowingStrategy.globalDefault())) .apply( ParDo.of(new ParDoMultiOverrideFactory.ToKeyedWorkItem<String, Integer>()) .withOutputTags(keyedTag, TupleTagList.empty())) .get(keyedTag) .setCoder( KeyedWorkItemCoder.of( StringUtf8Coder.of(), KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()), GlobalWindow.Coder.INSTANCE)); p.traverseTopologically(visitor); assertThat(visitor.getKeyedPValues(), hasItem(keyed)); }
@Test public void getKeyedPValuesBeforeTraverseThrows() { thrown.expect(IllegalStateException.class); thrown.expectMessage("completely traversed"); thrown.expectMessage("getKeyedPValues"); visitor.getKeyedPValues(); }