fix merge windows tests

Author: stankiewicz

runners/core-java/src/main/java/org/apache/beam/runners/core/LateDataDroppingDoFnRunner.java

@@ -145,7 +145,13 @@ public <K, InputT> Iterable<WindowedValue<InputT>> filter(
           } else {
             nonLateElements.add(
                 WindowedValues.of(
-                    element.getValue(), element.getTimestamp(), window, element.getPaneInfo()));
+                    element.getValue(),
+                    element.getTimestamp(),
+                    window,
+                    element.getPaneInfo(),
+                    element.getRecordId(),
+                    element.getRecordOffset(),
+                    element.causedByDrain()));
           }
         }
       }

runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnRunner.java

@@ -38,7 +38,7 @@
 import org.apache.beam.sdk.metrics.Counter;
 import org.apache.beam.sdk.metrics.Metrics;
 import org.apache.beam.sdk.options.PipelineOptions;
-import org.apache.beam.sdk.state.CombiningState;
+import org.apache.beam.sdk.state.BagState;
 import org.apache.beam.sdk.state.TimeDomain;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.transforms.windowing.PaneInfo;
@@ -108,11 +108,9 @@ public class ReduceFnRunner<K, InputT, OutputT, W extends BoundedWindow> {
    *   <li>It uses discarding or accumulation mode according to the {@link WindowingStrategy}.
    * </ul>
    */
-  static final StateTag<CombiningState<CombinedMetadata, CombinedMetadata, CombinedMetadata>>
-      METADATA_TAG =
-          StateTags.makeSystemTagInternal(
-              StateTags.combiningValue(
-                  "combinedMetadata", CombinedMetadata.Coder.of(), CombinedMetadataCombiner.of()));
+  static final StateTag<BagState<CombinedMetadata>> METADATA_TAG =
+      StateTags.makeSystemTagInternal(
+          StateTags.bag("combinedMetadata", CombinedMetadata.Coder.of()));
 
   private final WindowingStrategy<Object, W> windowingStrategy;
 
@@ -1010,6 +1008,7 @@ private boolean needToEmit(boolean isEmpty, boolean isFinished, PaneInfo.Timing
   private void prefetchOnTrigger(
       final ReduceFn<K, InputT, OutputT, W>.Context directContext,
       ReduceFn<K, InputT, OutputT, W>.Context renamedContext) {
+    directContext.state().access(METADATA_TAG).readLater();
     paneInfoTracker.prefetchPaneInfo(directContext);
     watermarkHold.prefetchExtract(renamedContext);
     nonEmptyPanes.isEmpty(renamedContext.state()).readLater();
@@ -1028,12 +1027,11 @@ private void prefetchOnTrigger(
       boolean isEndOfWindow,
       CombinedMetadata metadata)
       throws Exception {
-    CombiningState<CombinedMetadata, CombinedMetadata, CombinedMetadata> metadataState =
-        directContext.state().access(METADATA_TAG);
-    CombinedMetadata aggregatedMetadata = metadataState.read();
-    if (aggregatedMetadata == null) {
-      aggregatedMetadata = CombinedMetadata.createDefault();
-    }
+    BagState<CombinedMetadata> metadataState = directContext.state().access(METADATA_TAG);
+    Iterable<CombinedMetadata> allMetadata = metadataState.read();
+    CombinedMetadata aggregatedMetadata =
+        CombinedMetadataCombiner.of().mergeAccumulators(allMetadata);
+
     CombinedMetadata fullyAggregatedMetadata =
         CombinedMetadataCombiner.of().addInput(aggregatedMetadata, metadata);
     final CausedByDrain aggregatedCausedByDrain = fullyAggregatedMetadata.causedByDrain();

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java

@@ -39,6 +39,7 @@
 import org.apache.beam.sdk.transforms.windowing.PaneInfo;
 import org.apache.beam.sdk.transforms.windowing.PaneInfo.PaneInfoCoder;
 import org.apache.beam.sdk.util.ByteStringOutputStream;
+import org.apache.beam.sdk.values.CausedByDrain;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.ValueWithRecordId;
 import org.apache.beam.sdk.values.ValueWithRecordId.ValueWithRecordIdCoder;
@@ -220,7 +221,12 @@ public long add(WindowedValue<T> data) throws IOException {
       ByteString id = ByteString.EMPTY;
       // todo #33176 specify additional metadata in the future
       BeamFnApi.Elements.ElementMetadata additionalMetadata =
-          BeamFnApi.Elements.ElementMetadata.newBuilder().build();
+          BeamFnApi.Elements.ElementMetadata.newBuilder()
+              .setDrain(
+                  data.causedByDrain() == CausedByDrain.CAUSED_BY_DRAIN
+                      ? BeamFnApi.Elements.DrainMode.Enum.DRAINING
+                      : BeamFnApi.Elements.DrainMode.Enum.NOT_DRAINING)
+              .build();
       ByteString metadata =
           encodeMetadata(
               stream, windowsCoder, data.getWindows(), data.getPaneInfo(), additionalMetadata);

runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java

@@ -1922,6 +1922,11 @@ public void testMergeWindows() throws Exception {
         .getValueBuilder()
         .setTimestamp(0)
         .setData(ByteString.EMPTY);
+    dataBuilder
+        .addBagsBuilder()
+        .setTag(combinedMetadataTag)
+        .setStateFamily(stateFamily)
+        .addValues(ByteString.copyFrom(new byte[] {0x01, 0x02, 0x08, 0x01}));
     server.whenGetDataCalled().thenReturn(dataResponse.build());
 
     expectedBytesRead += dataBuilder.build().getSerializedSize();
@@ -2154,7 +2159,6 @@ public void testMergeWindowsCaching() throws Exception {
                 Windmill.TagBag.newBuilder()
                     .setTag(combinedMetadataTag)
                     .setStateFamily(stateFamily)
-                    .setDeleteAll(true)
                     .addValues(ByteString.copyFrom(new byte[] {0x01, 0x02, 0x08, 0x01}))
                     .build())));
 
@@ -2308,7 +2312,7 @@ public void testMergeWindowsCaching() throws Exception {
     CacheStats stats = worker.getStateCacheStats();
     LOG.info("cache stats {}", stats);
     assertEquals(1, stats.hitCount());
-    assertEquals(4, stats.missCount());
+    assertEquals(5, stats.missCount());
     worker.stop();
   }