Search in sources :

Example 1 with DirectGroupAlsoByWindow

use of org.apache.beam.runners.direct.DirectGroupByKey.DirectGroupAlsoByWindow in project beam by apache.

the class KeyedPValueTrackingVisitorTest method keyedInputWithKeyPreserving.

@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<String, WindowedValue<KV<String, Integer>>>()).apply(new DirectGroupAlsoByWindow<String, WindowedValue<KV<String, Integer>>>(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));
}
Also used : Instant(org.joda.time.Instant) TupleTag(org.apache.beam.sdk.values.TupleTag) KV(org.apache.beam.sdk.values.KV) KeyedWorkItem(org.apache.beam.runners.core.KeyedWorkItem) DirectGroupAlsoByWindow(org.apache.beam.runners.direct.DirectGroupByKey.DirectGroupAlsoByWindow) WindowedValue(org.apache.beam.sdk.util.WindowedValue) IntervalWindow(org.apache.beam.sdk.transforms.windowing.IntervalWindow) Test(org.junit.Test)

Aggregations

KeyedWorkItem (org.apache.beam.runners.core.KeyedWorkItem)1 DirectGroupAlsoByWindow (org.apache.beam.runners.direct.DirectGroupByKey.DirectGroupAlsoByWindow)1 IntervalWindow (org.apache.beam.sdk.transforms.windowing.IntervalWindow)1 WindowedValue (org.apache.beam.sdk.util.WindowedValue)1 KV (org.apache.beam.sdk.values.KV)1 TupleTag (org.apache.beam.sdk.values.TupleTag)1 Instant (org.joda.time.Instant)1 Test (org.junit.Test)1