Fixes some bugs with GroupIntoBatches.

sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupIntoBatches.java

@@ -492,9 +492,6 @@ public void processElement(
       storedBatchSize.readLater();
       minBufferedTs.readLater();
 
-      // Make sure we always include the current timestamp in the minBufferedTs.
-      minBufferedTs.add(elementTs.getMillis());
-
       LOG.debug("*** BATCH *** Add element for window {} ", window);
       if (shouldCareAboutWeight) {
         final long elementWeight = weigher.apply(element.getValue());
@@ -527,20 +524,22 @@ public void processElement(
         }
         storedBatchSizeBytes.add(elementWeight);
       }
+
+      long oldOutputTs =
+          MoreObjects.firstNonNull(
+              minBufferedTs.read(), BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis());
+
       batch.add(element.getValue());
       // Blind add is supported with combiningState
       storedBatchSize.add(1L);
-      // Add the timestamp back into minBufferedTs as it might be cleared by flushBatch above.
       minBufferedTs.add(elementTs.getMillis());
 
       final long num = storedBatchSize.read();
+      long currentMinBufferedTs = minBufferedTs.read();
 
       // If this is the first element in the batch or if the timer's output timestamp needs
       // modifying, then set a timer.
-      long oldOutputTs =
-          MoreObjects.firstNonNull(
-              minBufferedTs.read(), BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis());
-      boolean needsNewTimer = num == 1 || minBufferedTs.read() != oldOutputTs;
+      boolean needsNewTimer = num == 1 || currentMinBufferedTs < oldOutputTs;
       if (needsNewTimer) {
         if (shouldCareAboutMaxBufferingDuration) {
           long targetTs =
@@ -549,14 +548,14 @@ public void processElement(
                   bufferingTimer.getCurrentRelativeTime().getMillis()
                       + maxBufferingDuration.getMillis());
           bufferingTimer
-              .withOutputTimestamp(Instant.ofEpochMilli(minBufferedTs.read()))
+              .withOutputTimestamp(Instant.ofEpochMilli(currentMinBufferedTs))
               .set(Instant.ofEpochMilli(targetTs));
         } else {
           // The only way to hold the watermark is to set a timer. Since there is no buffering
           // timer, we set a dummy
           // timer at the end of the window to manage the hold.
           Instant windowEnd = window.maxTimestamp().plus(allowedLateness);
-          holdTimer.withOutputTimestamp(Instant.ofEpochMilli(minBufferedTs.read())).set(windowEnd);
+          holdTimer.withOutputTimestamp(Instant.ofEpochMilli(currentMinBufferedTs)).set(windowEnd);
         }
       }
 

sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/GroupIntoBatchesTest.java

@@ -805,4 +805,61 @@ private static <K> boolean checkBatchByteSizes(
     }
     return true;
   }
+
+  /**
+   * Tests that when elements arrive out of order, the watermark hold is updated to the minimum
+   * timestamp. With ofSize(3) and 2 elements, elements remain buffered while the watermark advances
+   * past the second element's timestamp. If the watermark hold were not updated to the minimum
+   * timestamp, the watermark would advance past the window of the second element and the batch
+   * would be dropped as late data downstream.
+   */
+  @Test
+  @Category({
+    ValidatesRunner.class,
+    NeedsRunner.class,
+    UsesTimersInParDo.class,
+    UsesStatefulParDo.class,
+    UsesTestStream.class,
+    UsesOnWindowExpiration.class
+  })
+  public void testOutOfOrderTimestampsInBatch() {
+    Instant baseTs = new Instant(0L);
+    TestStream<KV<String, String>> stream =
+        TestStream.create(KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()))
+            .advanceWatermarkTo(baseTs)
+            .addElements(
+                TimestampedValue.of(
+                    KV.of("key", "val1"), baseTs.plus(Duration.standardSeconds(10))))
+            .addElements(
+                TimestampedValue.of(KV.of("key", "val2"), baseTs.plus(Duration.standardSeconds(2))))
+            .advanceWatermarkTo(baseTs.plus(Duration.standardSeconds(7)))
+            .advanceWatermarkToInfinity();
+
+    PCollection<KV<String, Iterable<String>>> output =
+        pipeline
+            .apply(stream)
+            .apply(GroupIntoBatches.ofSize(3))
+            .apply(Window.into(FixedWindows.of(Duration.standardSeconds(5))))
+            .apply(
+                ParDo.of(
+                    new DoFn<KV<String, Iterable<String>>, KV<String, String>>() {
+                      @ProcessElement
+                      public void processElement(
+                          @Element KV<String, Iterable<String>> element,
+                          OutputReceiver<KV<String, String>> receiver) {
+                        receiver.output(KV.of(element.getKey(), "val"));
+                      }
+                    }))
+            .apply(GroupByKey.create());
+
+    PAssert.that(output)
+        .satisfies(
+            iterable -> {
+              assertTrue(
+                  "Expected output from GroupByKey but was dropped as late data",
+                  iterable.iterator().hasNext());
+              return null;
+            });
+    pipeline.run();
+  }
 }