StatefulParDoEvaluatorFactory implements TransformEvaluatorFactory { @Override public <T> TransformEvaluator<T> forApplication( AppliedPTransform<?, ?, ?> application, CommittedBundle<?> inputBundle) throws Exception { @SuppressWarnings({"unchecked", "rawtypes"}) TransformEvaluator<T> evaluator = (TransformEvaluator<T>) createEvaluator((AppliedPTransform) application, (CommittedBundle) inputBundle); return evaluator; } StatefulParDoEvaluatorFactory(EvaluationContext evaluationContext, PipelineOptions options); @Override TransformEvaluator<T> forApplication(
      AppliedPTransform<?, ?, ?> application, CommittedBundle<?> inputBundle); @Override void cleanup();  }
@Test public void testUnprocessedElements() throws Exception { final String stateId = "my-state-id"; PCollection<KV<String, Integer>> mainInput = pipeline .apply(Create.of(KV.of("hello", 1), KV.of("hello", 2))) .apply(Window.into(FixedWindows.of(Duration.millis(10)))); final PCollectionView<List<Integer>> sideInput = pipeline .apply("Create side input", Create.of(42)) .apply("Window side input", Window.into(FixedWindows.of(Duration.millis(10)))) .apply("View side input", View.asList()); TupleTag<Integer> mainOutput = new TupleTag<>(); PCollection<Integer> produced = mainInput .apply( new ParDoMultiOverrideFactory.GbkThenStatefulParDo<>( new DoFn<KV<String, Integer>, Integer>() { @StateId(stateId) private final StateSpec<ValueState<String>> spec = StateSpecs.value(StringUtf8Coder.of()); @ProcessElement public void process(ProcessContext c) {} }, mainOutput, TupleTagList.empty(), Collections.singletonList(sideInput), DoFnSchemaInformation.create(), Collections.emptyMap())) .get(mainOutput) .setCoder(VarIntCoder.of()); StatefulParDoEvaluatorFactory<String, Integer, Integer> factory = new StatefulParDoEvaluatorFactory<>(mockEvaluationContext, options); AppliedPTransform< PCollection<KeyedWorkItem<String, KV<String, Integer>>>, PCollectionTuple, StatefulParDo<String, Integer, Integer>> producingTransform = (AppliedPTransform) DirectGraphs.getProducer(produced); when(mockEvaluationContext.getExecutionContext( eq(producingTransform), Mockito.<StructuralKey>any())) .thenReturn(mockExecutionContext); when(mockExecutionContext.getStepContext(any())).thenReturn(mockStepContext); when(mockEvaluationContext.createBundle(Matchers.<PCollection<Integer>>any())) .thenReturn(mockUncommittedBundle); when(mockStepContext.getTimerUpdate()).thenReturn(TimerUpdate.empty()); when(mockEvaluationContext.createSideInputReader(anyList())).thenReturn(mockSideInputReader); when(mockSideInputReader.isReady(Matchers.any(), Matchers.any())).thenReturn(false); IntervalWindow firstWindow = new IntervalWindow(new Instant(0), new Instant(9)); String key = "hello"; WindowedValue<KV<String, Integer>> firstKv = WindowedValue.of(KV.of(key, 1), new Instant(3), firstWindow, PaneInfo.NO_FIRING); WindowedValue<KeyedWorkItem<String, KV<String, Integer>>> gbkOutputElement = firstKv.withValue( KeyedWorkItems.elementsWorkItem( "hello", ImmutableList.of( firstKv, firstKv.withValue(KV.of(key, 13)), firstKv.withValue(KV.of(key, 15))))); CommittedBundle<KeyedWorkItem<String, KV<String, Integer>>> inputBundle = BUNDLE_FACTORY .createBundle( (PCollection<KeyedWorkItem<String, KV<String, Integer>>>) Iterables.getOnlyElement( TransformInputs.nonAdditionalInputs(producingTransform))) .add(gbkOutputElement) .commit(Instant.now()); TransformEvaluator<KeyedWorkItem<String, KV<String, Integer>>> evaluator = factory.forApplication(producingTransform, inputBundle); evaluator.processElement(gbkOutputElement); TransformResult<KeyedWorkItem<String, KV<String, Integer>>> result = evaluator.finishBundle(); List<Integer> pushedBackInts = new ArrayList<>(); for (WindowedValue<? extends KeyedWorkItem<String, KV<String, Integer>>> unprocessedElement : result.getUnprocessedElements()) { assertThat( Iterables.getOnlyElement(unprocessedElement.getWindows()), equalTo((BoundedWindow) firstWindow)); assertThat(unprocessedElement.getValue().key(), equalTo("hello")); for (WindowedValue<KV<String, Integer>> windowedKv : unprocessedElement.getValue().elementsIterable()) { pushedBackInts.add(windowedKv.getValue().getValue()); } } assertThat(pushedBackInts, containsInAnyOrder(1, 13, 15)); }
@Test public void windowCleanupScheduled() throws Exception { final String stateId = "my-state-id"; PCollection<KV<String, Integer>> input = pipeline .apply(Create.of(KV.of("hello", 1), KV.of("hello", 2))) .apply(Window.into(FixedWindows.of(Duration.millis(10)))); TupleTag<Integer> mainOutput = new TupleTag<>(); final ParDoMultiOverrideFactory.GbkThenStatefulParDo<String, Integer, Integer> gbkThenStatefulParDo; gbkThenStatefulParDo = new ParDoMultiOverrideFactory.GbkThenStatefulParDo<>( new DoFn<KV<String, Integer>, Integer>() { @StateId(stateId) private final StateSpec<ValueState<String>> spec = StateSpecs.value(StringUtf8Coder.of()); @ProcessElement public void process(ProcessContext c) {} }, mainOutput, TupleTagList.empty(), Collections.emptyList(), DoFnSchemaInformation.create(), Collections.emptyMap()); final PCollection<KeyedWorkItem<String, KV<String, Integer>>> grouped = gbkThenStatefulParDo.groupToKeyedWorkItem(input); PCollection<Integer> produced = gbkThenStatefulParDo.applyStatefulParDo(grouped).get(mainOutput).setCoder(VarIntCoder.of()); StatefulParDoEvaluatorFactory<String, Integer, Integer> factory = new StatefulParDoEvaluatorFactory<>(mockEvaluationContext, options); AppliedPTransform< PCollection<? extends KeyedWorkItem<String, KV<String, Integer>>>, PCollectionTuple, StatefulParDo<String, Integer, Integer>> producingTransform = (AppliedPTransform) DirectGraphs.getProducer(produced); when(mockEvaluationContext.getExecutionContext( eq(producingTransform), Mockito.<StructuralKey>any())) .thenReturn(mockExecutionContext); when(mockExecutionContext.getStepContext(any())).thenReturn(mockStepContext); IntervalWindow firstWindow = new IntervalWindow(new Instant(0), new Instant(9)); IntervalWindow secondWindow = new IntervalWindow(new Instant(10), new Instant(19)); StateNamespace firstWindowNamespace = StateNamespaces.window(IntervalWindow.getCoder(), firstWindow); StateNamespace secondWindowNamespace = StateNamespaces.window(IntervalWindow.getCoder(), secondWindow); StateTag<ValueState<String>> tag = StateTags.tagForSpec(stateId, StateSpecs.value(StringUtf8Coder.of())); stateInternals.state(firstWindowNamespace, tag).write("first"); stateInternals.state(secondWindowNamespace, tag).write("second"); CommittedBundle<KeyedWorkItem<String, KV<String, Integer>>> inputBundle = BUNDLE_FACTORY .createBundle(grouped) .add( WindowedValue.of( KeyedWorkItems.<String, KV<String, Integer>>elementsWorkItem( "hello", Collections.singleton( WindowedValue.of( KV.of("hello", 1), new Instant(3), firstWindow, PaneInfo.NO_FIRING))), new Instant(3), firstWindow, PaneInfo.NO_FIRING)) .add( WindowedValue.of( KeyedWorkItems.<String, KV<String, Integer>>elementsWorkItem( "hello", Collections.singleton( WindowedValue.of( KV.of("hello", 2), new Instant(11), secondWindow, PaneInfo.NO_FIRING))), new Instant(11), secondWindow, PaneInfo.NO_FIRING)) .commit(Instant.now()); factory.forApplication(producingTransform, inputBundle); ArgumentCaptor<Runnable> argumentCaptor = ArgumentCaptor.forClass(Runnable.class); verify(mockEvaluationContext) .scheduleAfterWindowExpiration( eq(producingTransform), eq(firstWindow), Mockito.any(), argumentCaptor.capture()); argumentCaptor.getValue().run(); assertThat(stateInternals.state(firstWindowNamespace, tag).read(), nullValue()); assertThat(stateInternals.state(secondWindowNamespace, tag).read(), equalTo("second")); verify(mockEvaluationContext) .scheduleAfterWindowExpiration( eq(producingTransform), eq(secondWindow), Mockito.any(), argumentCaptor.capture()); argumentCaptor.getValue().run(); assertThat(stateInternals.state(secondWindowNamespace, tag).read(), nullValue()); }