Code review comment: We can do a direct assignment here

Code before review:
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;
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());
}

Line under review: @StateId(stateId)

Code after addressing the review:
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());
}