UserParDoFnFactory implements ParDoFnFactory { @Override public ParDoFn create( PipelineOptions options, CloudObject cloudUserFn, @Nullable List<SideInputInfo> sideInputInfos, TupleTag<?> mainOutputTag, Map<TupleTag<?>, Integer> outputTupleTagsToReceiverIndices, DataflowExecutionContext<?> executionContext, DataflowOperationContext operationContext) throws Exception { DoFnInstanceManager instanceManager = fnCache.get( operationContext.nameContext().systemName(), () -> DoFnInstanceManagers.cloningPool(doFnExtractor.getDoFnInfo(cloudUserFn))); DoFnInfo<?, ?> doFnInfo = instanceManager.peek(); DataflowExecutionContext.DataflowStepContext stepContext = executionContext.getStepContext(operationContext); Iterable<PCollectionView<?>> sideInputViews = doFnInfo.getSideInputViews(); SideInputReader sideInputReader = executionContext.getSideInputReader(sideInputInfos, sideInputViews, operationContext); if (doFnInfo.getDoFn() instanceof BatchStatefulParDoOverrides.BatchStatefulDoFn) { BatchStatefulParDoOverrides.BatchStatefulDoFn fn = (BatchStatefulParDoOverrides.BatchStatefulDoFn) doFnInfo.getDoFn(); DoFn underlyingFn = fn.getUnderlyingDoFn(); return new BatchModeUngroupingParDoFn( (BatchModeExecutionContext.StepContext) stepContext, new SimpleParDoFn( options, DoFnInstanceManagers.singleInstance(doFnInfo.withFn(underlyingFn)), sideInputReader, doFnInfo.getMainOutput(), outputTupleTagsToReceiverIndices, stepContext, operationContext, doFnInfo.getDoFnSchemaInformation(), doFnInfo.getSideInputMapping(), runnerFactory)); } else if (doFnInfo.getDoFn() instanceof StreamingPCollectionViewWriterFn) { checkArgument( stepContext instanceof StreamingModeExecutionContext.StreamingModeStepContext, "stepContext must be a StreamingModeStepContext to use StreamingPCollectionViewWriterFn"); DataflowRunner.StreamingPCollectionViewWriterFn<Object> writerFn = (StreamingPCollectionViewWriterFn<Object>) doFnInfo.getDoFn(); return new StreamingPCollectionViewWriterParDoFn( (StreamingModeExecutionContext.StreamingModeStepContext) stepContext, writerFn.getView().getTagInternal(), writerFn.getDataCoder(), (Coder<BoundedWindow>) doFnInfo.getWindowingStrategy().getWindowFn().windowCoder()); } else { return new SimpleParDoFn( options, instanceManager, sideInputReader, doFnInfo.getMainOutput(), outputTupleTagsToReceiverIndices, stepContext, operationContext, doFnInfo.getDoFnSchemaInformation(), doFnInfo.getSideInputMapping(), runnerFactory); } } UserParDoFnFactory(DoFnExtractor doFnExtractor, DoFnRunnerFactory runnerFactory); @Override ParDoFn create(
      PipelineOptions options,
      CloudObject cloudUserFn,
      @Nullable List<SideInputInfo> sideInputInfos,
      TupleTag<?> mainOutputTag,
      Map<TupleTag<?>, Integer> outputTupleTagsToReceiverIndices,
      DataflowExecutionContext<?> executionContext,
      DataflowOperationContext operationContext);  }
@Test public void testFactoryReuseInStep() throws Exception { PipelineOptions options = PipelineOptionsFactory.create(); CounterSet counters = new CounterSet(); TestDoFn initialFn = new TestDoFn(Collections.<TupleTag<String>>emptyList()); CloudObject cloudObject = getCloudObject(initialFn); TestOperationContext operationContext = TestOperationContext.create(counters); ParDoFn parDoFn = factory.create( options, cloudObject, null, MAIN_OUTPUT, ImmutableMap.<TupleTag<?>, Integer>of(MAIN_OUTPUT, 0), BatchModeExecutionContext.forTesting(options, "testStage"), operationContext); Receiver rcvr = new OutputReceiver(); parDoFn.startBundle(rcvr); parDoFn.processElement(WindowedValue.valueInGlobalWindow("foo")); TestDoFn fn = (TestDoFn) ((SimpleParDoFn) parDoFn).getDoFnInfo().getDoFn(); assertThat(fn, not(theInstance(initialFn))); parDoFn.finishBundle(); assertThat(fn.state, equalTo(TestDoFn.State.FINISHED)); ParDoFn secondParDoFn = factory.create( options, cloudObject, null, MAIN_OUTPUT, ImmutableMap.<TupleTag<?>, Integer>of(MAIN_OUTPUT, 0), BatchModeExecutionContext.forTesting(options, "testStage"), operationContext); assertThat(fn.state, equalTo(TestDoFn.State.FINISHED)); secondParDoFn.startBundle(rcvr); secondParDoFn.processElement(WindowedValue.valueInGlobalWindow("spam")); TestDoFn reobtainedFn = (TestDoFn) ((SimpleParDoFn) secondParDoFn).getDoFnInfo().getDoFn(); secondParDoFn.finishBundle(); assertThat(reobtainedFn.state, equalTo(TestDoFn.State.FINISHED)); assertThat(fn, theInstance(reobtainedFn)); }
@Test public void testFactorySimultaneousUse() throws Exception { PipelineOptions options = PipelineOptionsFactory.create(); CounterSet counters = new CounterSet(); TestDoFn initialFn = new TestDoFn(Collections.<TupleTag<String>>emptyList()); CloudObject cloudObject = getCloudObject(initialFn); ParDoFn parDoFn = factory.create( options, cloudObject, null, MAIN_OUTPUT, ImmutableMap.<TupleTag<?>, Integer>of(MAIN_OUTPUT, 0), BatchModeExecutionContext.forTesting(options, "testStage"), TestOperationContext.create(counters)); ParDoFn secondParDoFn = factory.create( options, cloudObject, null, MAIN_OUTPUT, ImmutableMap.<TupleTag<?>, Integer>of(MAIN_OUTPUT, 0), BatchModeExecutionContext.forTesting(options, "testStage"), TestOperationContext.create(counters)); Receiver rcvr = new OutputReceiver(); parDoFn.startBundle(rcvr); parDoFn.processElement(WindowedValue.valueInGlobalWindow("foo")); TestDoFn firstDoFn = (TestDoFn) ((SimpleParDoFn) parDoFn).getDoFnInfo().getDoFn(); secondParDoFn.startBundle(rcvr); secondParDoFn.processElement(WindowedValue.valueInGlobalWindow("spam")); TestDoFn secondDoFn = (TestDoFn) ((SimpleParDoFn) secondParDoFn).getDoFnInfo().getDoFn(); parDoFn.finishBundle(); secondParDoFn.finishBundle(); assertThat(firstDoFn, not(theInstance(secondDoFn))); assertThat(firstDoFn.state, equalTo(TestDoFn.State.FINISHED)); assertThat(secondDoFn.state, equalTo(TestDoFn.State.FINISHED)); }
@Test public void testFactoryDoesNotReuseAfterAborted() throws Exception { PipelineOptions options = PipelineOptionsFactory.create(); CounterSet counters = new CounterSet(); TestDoFn initialFn = new TestDoFn(Collections.<TupleTag<String>>emptyList()); CloudObject cloudObject = getCloudObject(initialFn); ParDoFn parDoFn = factory.create( options, cloudObject, null, MAIN_OUTPUT, ImmutableMap.<TupleTag<?>, Integer>of(MAIN_OUTPUT, 0), BatchModeExecutionContext.forTesting(options, "testStage"), TestOperationContext.create(counters)); Receiver rcvr = new OutputReceiver(); parDoFn.startBundle(rcvr); parDoFn.processElement(WindowedValue.valueInGlobalWindow("foo")); TestDoFn fn = (TestDoFn) ((SimpleParDoFn) parDoFn).getDoFnInfo().getDoFn(); parDoFn.abort(); assertThat(fn.state, equalTo(TestDoFn.State.TORN_DOWN)); ParDoFn secondParDoFn = factory.create( options, cloudObject.clone(), null, MAIN_OUTPUT, ImmutableMap.<TupleTag<?>, Integer>of(MAIN_OUTPUT, 0), BatchModeExecutionContext.forTesting(options, "testStage"), TestOperationContext.create(counters)); secondParDoFn.startBundle(rcvr); secondParDoFn.processElement(WindowedValue.valueInGlobalWindow("foo")); TestDoFn secondFn = (TestDoFn) ((SimpleParDoFn) secondParDoFn).getDoFnInfo().getDoFn(); assertThat(secondFn, not(theInstance(fn))); assertThat(fn.state, equalTo(TestDoFn.State.TORN_DOWN)); assertThat(secondFn.state, equalTo(TestDoFn.State.PROCESSING)); }
@Test public void testCleanupRegistered() throws Exception { PipelineOptions options = PipelineOptionsFactory.create(); CounterSet counters = new CounterSet(); DoFn<?, ?> initialFn = new TestStatefulDoFn(); CloudObject cloudObject = getCloudObject( initialFn, WindowingStrategy.globalDefault().withWindowFn(FixedWindows.of(Duration.millis(10)))); TimerInternals timerInternals = mock(TimerInternals.class); DataflowStepContext stepContext = mock(DataflowStepContext.class); when(stepContext.timerInternals()).thenReturn(timerInternals); DataflowExecutionContext<DataflowStepContext> executionContext = mock(DataflowExecutionContext.class); TestOperationContext operationContext = TestOperationContext.create(counters); when(executionContext.getStepContext(operationContext)).thenReturn(stepContext); when(executionContext.getSideInputReader(any(), any(), any())) .thenReturn(NullSideInputReader.empty()); ParDoFn parDoFn = factory.create( options, cloudObject, Collections.emptyList(), MAIN_OUTPUT, ImmutableMap.of(MAIN_OUTPUT, 0), executionContext, operationContext); Receiver rcvr = new OutputReceiver(); parDoFn.startBundle(rcvr); IntervalWindow firstWindow = new IntervalWindow(new Instant(0), new Instant(10)); parDoFn.processElement( WindowedValue.of("foo", new Instant(1), firstWindow, PaneInfo.NO_FIRING)); verify(stepContext) .setStateCleanupTimer( SimpleParDoFn.CLEANUP_TIMER_ID, firstWindow, IntervalWindow.getCoder(), firstWindow.maxTimestamp().plus(1L), firstWindow.maxTimestamp().plus(1L)); }
@Test public void testCleanupWorks() throws Exception { PipelineOptions options = PipelineOptionsFactory.create(); CounterSet counters = new CounterSet(); DoFn<?, ?> initialFn = new TestStatefulDoFn(); CloudObject cloudObject = getCloudObject(initialFn, WindowingStrategy.of(FixedWindows.of(Duration.millis(10)))); StateInternals stateInternals = InMemoryStateInternals.forKey("dummy"); DataflowStepContext stepContext = mock(DataflowStepContext.class); DataflowStepContext userStepContext = mock(DataflowStepContext.class); when(stepContext.namespacedToUser()).thenReturn(userStepContext); when(stepContext.stateInternals()).thenReturn(stateInternals); when(userStepContext.stateInternals()).thenReturn((StateInternals) stateInternals); DataflowExecutionContext<DataflowStepContext> executionContext = mock(DataflowExecutionContext.class); TestOperationContext operationContext = TestOperationContext.create(counters); when(executionContext.getStepContext(operationContext)).thenReturn(stepContext); when(executionContext.getSideInputReader(any(), any(), any())) .thenReturn(NullSideInputReader.empty()); ParDoFn parDoFn = factory.create( options, cloudObject, Collections.emptyList(), MAIN_OUTPUT, ImmutableMap.of(MAIN_OUTPUT, 0), executionContext, operationContext); Receiver rcvr = new OutputReceiver(); parDoFn.startBundle(rcvr); IntervalWindow firstWindow = new IntervalWindow(new Instant(0), new Instant(9)); IntervalWindow secondWindow = new IntervalWindow(new Instant(10), new Instant(19)); Coder<IntervalWindow> windowCoder = IntervalWindow.getCoder(); StateNamespace firstWindowNamespace = StateNamespaces.window(windowCoder, firstWindow); StateNamespace secondWindowNamespace = StateNamespaces.window(windowCoder, secondWindow); StateTag<ValueState<String>> tag = StateTags.tagForSpec(TestStatefulDoFn.STATE_ID, StateSpecs.value(StringUtf8Coder.of())); stateInternals.state(firstWindowNamespace, tag).write("first"); stateInternals.state(secondWindowNamespace, tag).write("second"); when(userStepContext.getNextFiredTimer(windowCoder)).thenReturn(null); when(stepContext.getNextFiredTimer(windowCoder)) .thenReturn( TimerData.of( SimpleParDoFn.CLEANUP_TIMER_ID, firstWindowNamespace, firstWindow.maxTimestamp().plus(1L), firstWindow.maxTimestamp().plus(1L), TimeDomain.EVENT_TIME)) .thenReturn(null); parDoFn.processTimers(); assertThat(stateInternals.state(firstWindowNamespace, tag).read(), nullValue()); assertThat(stateInternals.state(secondWindowNamespace, tag).read(), equalTo("second")); when(stepContext.getNextFiredTimer((Coder) windowCoder)) .thenReturn( TimerData.of( SimpleParDoFn.CLEANUP_TIMER_ID, secondWindowNamespace, secondWindow.maxTimestamp().plus(1L), secondWindow.maxTimestamp().plus(1L), TimeDomain.EVENT_TIME)) .thenReturn(null); parDoFn.processTimers(); assertThat(stateInternals.state(firstWindowNamespace, tag).read(), nullValue()); assertThat(stateInternals.state(secondWindowNamespace, tag).read(), nullValue()); }
@Test public void testCleanupRegistered() throws Exception { PipelineOptions options = PipelineOptionsFactory.create(); CounterSet counters = new CounterSet(); DoFn<?, ?> initialFn = new TestStatefulDoFn(); CloudObject cloudObject = getCloudObject( initialFn, WindowingStrategy.globalDefault().withWindowFn(FixedWindows.of(Duration.millis(10)))); TimerInternals timerInternals = mock(TimerInternals.class); DataflowStepContext stepContext = mock(DataflowStepContext.class); when(stepContext.timerInternals()).thenReturn(timerInternals); DataflowExecutionContext<DataflowStepContext> executionContext = mock(DataflowExecutionContext.class); TestOperationContext operationContext = TestOperationContext.create(counters); when(executionContext.getStepContext(operationContext)).thenReturn(stepContext); when(executionContext.getSideInputReader(any(), any(), any())) .thenReturn(NullSideInputReader.empty()); ParDoFn parDoFn = factory.create( options, cloudObject, Collections.emptyList(), MAIN_OUTPUT, ImmutableMap.of(MAIN_OUTPUT, 0), executionContext, operationContext); Receiver rcvr = new OutputReceiver(); parDoFn.startBundle(rcvr); IntervalWindow firstWindow = new IntervalWindow(new Instant(0), new Instant(10)); parDoFn.processElement( WindowedValue.of("foo", new Instant(1), firstWindow, PaneInfo.NO_FIRING)); verify(stepContext) .setStateCleanupTimer( SimpleParDoFn.CLEANUP_TIMER_ID, firstWindow, IntervalWindow.getCoder(), firstWindow.maxTimestamp().plus(1L)); }
@Test public void testCleanupWorks() throws Exception { PipelineOptions options = PipelineOptionsFactory.create(); CounterSet counters = new CounterSet(); DoFn<?, ?> initialFn = new TestStatefulDoFn(); CloudObject cloudObject = getCloudObject(initialFn, WindowingStrategy.of(FixedWindows.of(Duration.millis(10)))); StateInternals stateInternals = InMemoryStateInternals.forKey("dummy"); DataflowStepContext stepContext = mock(DataflowStepContext.class); DataflowStepContext userStepContext = mock(DataflowStepContext.class); when(stepContext.namespacedToUser()).thenReturn(userStepContext); when(userStepContext.stateInternals()).thenReturn((StateInternals) stateInternals); DataflowExecutionContext<DataflowStepContext> executionContext = mock(DataflowExecutionContext.class); TestOperationContext operationContext = TestOperationContext.create(counters); when(executionContext.getStepContext(operationContext)).thenReturn(stepContext); when(executionContext.getSideInputReader(any(), any(), any())) .thenReturn(NullSideInputReader.empty()); ParDoFn parDoFn = factory.create( options, cloudObject, Collections.emptyList(), MAIN_OUTPUT, ImmutableMap.of(MAIN_OUTPUT, 0), executionContext, operationContext); Receiver rcvr = new OutputReceiver(); parDoFn.startBundle(rcvr); IntervalWindow firstWindow = new IntervalWindow(new Instant(0), new Instant(9)); IntervalWindow secondWindow = new IntervalWindow(new Instant(10), new Instant(19)); Coder<IntervalWindow> windowCoder = IntervalWindow.getCoder(); StateNamespace firstWindowNamespace = StateNamespaces.window(windowCoder, firstWindow); StateNamespace secondWindowNamespace = StateNamespaces.window(windowCoder, secondWindow); StateTag<ValueState<String>> tag = StateTags.tagForSpec(TestStatefulDoFn.STATE_ID, StateSpecs.value(StringUtf8Coder.of())); stateInternals.state(firstWindowNamespace, tag).write("first"); stateInternals.state(secondWindowNamespace, tag).write("second"); when(userStepContext.getNextFiredTimer(windowCoder)).thenReturn(null); when(stepContext.getNextFiredTimer(windowCoder)) .thenReturn( TimerData.of( SimpleParDoFn.CLEANUP_TIMER_ID, firstWindowNamespace, firstWindow.maxTimestamp().plus(1L), TimeDomain.EVENT_TIME)) .thenReturn(null); parDoFn.processTimers(); assertThat(stateInternals.state(firstWindowNamespace, tag).read(), nullValue()); assertThat(stateInternals.state(secondWindowNamespace, tag).read(), equalTo("second")); when(stepContext.getNextFiredTimer((Coder) windowCoder)) .thenReturn( TimerData.of( SimpleParDoFn.CLEANUP_TIMER_ID, secondWindowNamespace, secondWindow.maxTimestamp().plus(1L), TimeDomain.EVENT_TIME)) .thenReturn(null); parDoFn.processTimers(); assertThat(stateInternals.state(firstWindowNamespace, tag).read(), nullValue()); assertThat(stateInternals.state(secondWindowNamespace, tag).read(), nullValue()); }