Fixes some bugs with GroupIntoBatches.

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

@@ -43,7 +43,6 @@
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects;
-import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
 import org.joda.time.Duration;
 import org.joda.time.Instant;
@@ -78,13 +77,13 @@
  * PCollection<KV<String, Iterable<String>>> batched = input
  *     .apply(GroupIntoBatches.<String, String>ofSize(batchSize))
  *     .setCoder(KvCoder.of(StringUtf8Coder.of(), IterableCoder.of(StringUtf8Coder.of())))
- *     .apply(ParDo.of(new DoFn<KV<String, Iterable<String>>, KV<String, String>>() }{
- *        {@code @ProcessElement
+ *     .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>> r) {
  *             r.output(KV.of(element.getKey(), callWebService(element.getValue())));
  *         }
- *     }}));
+ *     }));
  * </pre>
  *
  * <p>Example 2 (batch unbounded input in a global window):
@@ -171,7 +170,7 @@ private GroupIntoBatches(BatchingParams<InputT> params) {
 
   /** Aim to create batches each with the specified element count. */
   public static <K, InputT> GroupIntoBatches<K, InputT> ofSize(long batchSize) {
-    Preconditions.checkState(batchSize < Long.MAX_VALUE);
+    checkArgument(batchSize < Long.MAX_VALUE, "batchSize must be less than Long.MAX_VALUE");
     return new GroupIntoBatches<K, InputT>(BatchingParams.createDefault()).withSize(batchSize);
   }
 
@@ -210,9 +209,23 @@ public String toString() {
     return super.toString() + ", params=" + params;
   }
 
+  @Override
+  public void populateDisplayData(DisplayData.Builder builder) {
+    super.populateDisplayData(builder);
+    if (params.getBatchSize() < Long.MAX_VALUE) {
+      builder.add(DisplayData.item("batchSize", params.getBatchSize()));
+    }
+    if (params.getBatchSizeBytes() < Long.MAX_VALUE) {
+      builder.add(DisplayData.item("batchSizeBytes", params.getBatchSizeBytes()));
+    }
+    if (params.getMaxBufferingDuration().isLongerThan(Duration.ZERO)) {
+      builder.add(DisplayData.item("maxBufferingDuration", params.getMaxBufferingDuration()));
+    }
+  }
+
   /** See {@link #ofSize(long)}. */
   public GroupIntoBatches<K, InputT> withSize(long batchSize) {
-    Preconditions.checkState(batchSize < Long.MAX_VALUE);
+    checkArgument(batchSize < Long.MAX_VALUE, "batchSize must be less than Long.MAX_VALUE");
     return new GroupIntoBatches<>(
         BatchingParams.create(
             batchSize,
@@ -223,7 +236,8 @@ public GroupIntoBatches<K, InputT> withSize(long batchSize) {
 
   /** See {@link #ofByteSize(long)}. */
   public GroupIntoBatches<K, InputT> withByteSize(long batchSizeBytes) {
-    Preconditions.checkState(batchSizeBytes < Long.MAX_VALUE);
+    checkArgument(
+        batchSizeBytes < Long.MAX_VALUE, "batchSizeBytes must be less than Long.MAX_VALUE");
     return new GroupIntoBatches<>(
         BatchingParams.create(
             params.getBatchSize(),
@@ -235,7 +249,8 @@ public GroupIntoBatches<K, InputT> withByteSize(long batchSizeBytes) {
   /** See {@link #ofByteSize(long, SerializableFunction)}. */
   public GroupIntoBatches<K, InputT> withByteSize(
       long batchSizeBytes, SerializableFunction<InputT, Long> getElementByteSize) {
-    Preconditions.checkState(batchSizeBytes < Long.MAX_VALUE);
+    checkArgument(
+        batchSizeBytes < Long.MAX_VALUE, "batchSizeBytes must be less than Long.MAX_VALUE");
     return new GroupIntoBatches<>(
         BatchingParams.create(
             params.getBatchSize(),
@@ -308,6 +323,20 @@ public KV<ShardedKey<K>, InputT> apply(KV<K, InputT> input) {
           .setCoder(KvCoder.of(ShardedKey.Coder.of(keyCoder), valueCoder))
           .apply(new GroupIntoBatches<>(getBatchingParams()));
     }
+
+    @Override
+    public void populateDisplayData(DisplayData.Builder builder) {
+      super.populateDisplayData(builder);
+      if (params.getBatchSize() < Long.MAX_VALUE) {
+        builder.add(DisplayData.item("batchSize", params.getBatchSize()));
+      }
+      if (params.getBatchSizeBytes() < Long.MAX_VALUE) {
+        builder.add(DisplayData.item("batchSizeBytes", params.getBatchSizeBytes()));
+      }
+      if (params.getMaxBufferingDuration().isLongerThan(Duration.ZERO)) {
+        builder.add(DisplayData.item("maxBufferingDuration", params.getMaxBufferingDuration()));
+      }
+    }
   }
 
   private static class ByteSizeObserver extends ElementByteSizeObserver {
@@ -322,7 +351,6 @@ public long getElementByteSize() {
       return this.elementByteSize;
     }
   }
-  ;
 
   @Override
   public PCollection<KV<K, Iterable<InputT>>> expand(PCollection<KV<K, InputT>> input) {
@@ -492,13 +520,11 @@ 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());
-        if (elementWeight + storedBatchSizeBytes.read() > batchSizeBytes) {
+        if (storedBatchSizeBytes.read() > 0
+            && elementWeight + storedBatchSizeBytes.read() > batchSizeBytes) {
           // Firing by count and size limits behave differently.
           //
           // We always increase the count by one, so we will fire at the exact limit without any
@@ -527,36 +553,47 @@ public void processElement(
         }
         storedBatchSizeBytes.add(elementWeight);
       }
+
+      long oldMinBufferedTs =
+          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 < oldMinBufferedTs;
       if (needsNewTimer) {
         if (shouldCareAboutMaxBufferingDuration) {
-          long targetTs =
-              MoreObjects.firstNonNull(
-                  timerTs.read(),
-                  bufferingTimer.getCurrentRelativeTime().getMillis()
-                      + maxBufferingDuration.getMillis());
+          Long existingTargetTs = timerTs.read();
+          long targetTs;
+          if (existingTargetTs == null) {
+            targetTs =
+                bufferingTimer.getCurrentRelativeTime().getMillis()
+                    + maxBufferingDuration.getMillis();
+            timerTs.write(targetTs);
+          } else {
+            targetTs = existingTargetTs;
+          }
           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);
+          Instant windowMax = window.maxTimestamp();
+          Instant windowEnd =
+              windowMax.isBefore(BoundedWindow.TIMESTAMP_MAX_VALUE.minus(allowedLateness))
+                  ? windowMax.plus(allowedLateness)
+                  : BoundedWindow.TIMESTAMP_MAX_VALUE;
+          holdTimer.withOutputTimestamp(Instant.ofEpochMilli(currentMinBufferedTs)).set(windowEnd);
         }
       }
 

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

@@ -17,6 +17,8 @@
  */
 package org.apache.beam.sdk.transforms;
 
+import static org.apache.beam.sdk.transforms.display.DisplayDataMatchers.hasDisplayItem;
+import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 
@@ -44,6 +46,7 @@
 import org.apache.beam.sdk.testing.UsesTestStreamWithProcessingTime;
 import org.apache.beam.sdk.testing.UsesTimersInParDo;
 import org.apache.beam.sdk.testing.ValidatesRunner;
+import org.apache.beam.sdk.transforms.display.DisplayData;
 import org.apache.beam.sdk.transforms.windowing.AfterPane;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.transforms.windowing.FixedWindows;
@@ -805,4 +808,71 @@ private static <K> boolean checkBatchByteSizes(
     }
     return true;
   }
+
+  @Test(expected = IllegalArgumentException.class)
+  public void testMaxBatchSizeThrows() {
+    GroupIntoBatches.ofSize(Long.MAX_VALUE);
+  }
+
+  @Test(expected = IllegalArgumentException.class)
+  public void testMaxBatchSizeBytesThrows() {
+    GroupIntoBatches.ofByteSize(Long.MAX_VALUE);
+  }
+
+  @Test
+  public void testDisplayData() {
+    GroupIntoBatches<String, String> gib =
+        GroupIntoBatches.<String, String>ofSize(100)
+            .withByteSize(1000)
+            .withMaxBufferingDuration(Duration.standardSeconds(5));
+    DisplayData displayData = DisplayData.from(gib);
+    assertThat(displayData, hasDisplayItem("batchSize", 100L));
+    assertThat(displayData, hasDisplayItem("batchSizeBytes", 1000L));
+    assertThat(displayData, hasDisplayItem("maxBufferingDuration", Duration.standardSeconds(5)));
+  }
+
+  /**
+   * 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 until window expiration,
+   * exercising the updated timer watermark hold rather than flushing immediately in processElement.
+   */
+  @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(5))))
+            .advanceWatermarkToInfinity();
+
+    PCollection<Long> outputTimestamps =
+        pipeline
+            .apply(stream)
+            .apply(GroupIntoBatches.ofSize(3))
+            .apply(
+                ParDo.of(
+                    new DoFn<KV<String, Iterable<String>>, Long>() {
+                      @ProcessElement
+                      public void processElement(
+                          @Timestamp Instant ts, OutputReceiver<Long> receiver) {
+                        receiver.output(ts.getMillis());
+                      }
+                    }));
+
+    PAssert.that(outputTimestamps)
+        .containsInAnyOrder(baseTs.plus(Duration.standardSeconds(5)).getMillis());
+    pipeline.run();
+  }
 }