diff --git a/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/FlinkGroupAlsoByWindowWrapper.java b/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/FlinkGroupAlsoByWindowWrapper.java
index b413d7a595a4..3dc5a799f557 100644
--- a/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/FlinkGroupAlsoByWindowWrapper.java
+++ b/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/FlinkGroupAlsoByWindowWrapper.java
@@ -220,6 +220,7 @@ private FlinkGroupAlsoByWindowWrapper(PipelineOptions options,
   public void open() throws Exception {
     super.open();
     this.context = new ProcessContext(operator, new TimestampedCollector<>(output), this.timerInternals);
+    operator.startBundle(context);
   }
 
   /**
@@ -252,19 +253,23 @@ private <W extends BoundedWindow> DoFn<KeyedWorkItem<K, VIN>, KV<K, VOUT>> creat
 
   private void processKeyedWorkItem(KeyedWorkItem<K, VIN> workItem) throws Exception {
     context.setElement(workItem, getStateInternalsForKey(workItem.key()));
-
-    // TODO: Ideally startBundle/finishBundle would be called when the operator is first used / about to be discarded.
-    operator.startBundle(context);
     operator.processElement(context);
-    operator.finishBundle(context);
   }
 
   @Override
   public void processElement(StreamRecord<WindowedValue<KV<K, VIN>>> element) throws Exception {
-    ArrayList<WindowedValue<VIN>> elements = new ArrayList<>();
-    elements.add(WindowedValue.of(element.getValue().getValue().getValue(), element.getValue().getTimestamp(),
-        element.getValue().getWindows(), element.getValue().getPane()));
-    processKeyedWorkItem(KeyedWorkItems.elementsWorkItem(element.getValue().getValue().getKey(), elements));
+    final WindowedValue<KV<K, VIN>> windowedValue = element.getValue();
+    final KV<K, VIN> kv = windowedValue.getValue();
+
+    final WindowedValue<VIN> updatedWindowedValue = WindowedValue.of(kv.getValue(),
+        windowedValue.getTimestamp(),
+        windowedValue.getWindows(),
+        windowedValue.getPane());
+
+    processKeyedWorkItem(
+        KeyedWorkItems.elementsWorkItem(
+            kv.getKey(),
+            Collections.singletonList(updatedWindowedValue)));
   }
 
   @Override
@@ -309,6 +314,7 @@ public void processWatermark(Watermark mark) throws Exception {
 
   @Override
   public void close() throws Exception {
+    operator.finishBundle(context);
     super.close();
   }