boyuanzz commented on a change in pull request #12726:
URL: https://github.com/apache/beam/pull/12726#discussion_r487313819



##########
File path: 
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupIntoBatches.java
##########
@@ -66,25 +68,49 @@
  *         }
  *     }}));
  * </pre>
+ *
+ * <p>Example 2 (batch unbounded input in a global window):
+ *
+ * <pre>{@code
+ * PCollection<KV<String, String>> unboundedInput = ...;
+ * long batchSize = 100L;
+ * Duration maxBufferingDuration = Duration.standardSeconds(10);
+ * PCollection<KV<String, Iterable<String>>> batched = unboundedInput
+ *     .apply(Window.<KV<String, String>>into(new GlobalWindows())
+ *         .triggering(Repeatedly.forever(AfterPane.elementCountAtLeast(1)))
+ *         .discardingFiredPanes())
+ *     .apply(GroupIntoBatches.<String, String>ofSize(batchSize)
+ *         .withMaxBufferingDuration(maxBufferingDuration));
+ * }</pre>
  */
 public class GroupIntoBatches<K, InputT>
     extends PTransform<PCollection<KV<K, InputT>>, PCollection<KV<K, 
Iterable<InputT>>>> {
 
   private final long batchSize;
+  private final Duration maxBufferingDuration;
 
-  private GroupIntoBatches(long batchSize) {
+  private GroupIntoBatches(long batchSize, Duration maxBufferingDuration) {
     this.batchSize = batchSize;
+    this.maxBufferingDuration = maxBufferingDuration;
   }
 
   public static <K, InputT> GroupIntoBatches<K, InputT> ofSize(long batchSize) 
{
-    return new GroupIntoBatches<>(batchSize);
+    return new GroupIntoBatches<>(batchSize, Duration.ZERO);

Review comment:
       Can we use `null` here?

##########
File path: 
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupIntoBatches.java
##########
@@ -150,70 +185,96 @@ public long apply(long left, long right) {
               });
 
       this.keySpec = StateSpecs.value(inputKeyCoder);
-      // prefetch every 20% of batchSize elements. Do not prefetch if 
batchSize is too little
+      // Prefetch every 20% of batchSize elements. Do not prefetch if 
batchSize is too little
       this.prefetchFrequency = ((batchSize / 5) <= 1) ? Long.MAX_VALUE : 
(batchSize / 5);
     }
 
     @ProcessElement
     public void processElement(
-        @TimerId(END_OF_WINDOW_ID) Timer timer,
+        @TimerId(END_OF_WINDOW_ID) Timer windowTimer,
+        @TimerId(END_OF_BUFFERING_ID) Timer bufferingTimer,
         @StateId(BATCH_ID) BagState<InputT> batch,
         @StateId(NUM_ELEMENTS_IN_BATCH_ID) CombiningState<Long, long[], Long> 
numElementsInBatch,
         @StateId(KEY_ID) ValueState<K> key,
         @Element KV<K, InputT> element,
         BoundedWindow window,
         OutputReceiver<KV<K, Iterable<InputT>>> receiver) {
-      Instant windowExpires = window.maxTimestamp().plus(allowedLateness);
-
-      LOG.debug(
-          "*** SET TIMER *** to point in time {} for window {}",
-          windowExpires.toString(),
-          window.toString());
-      timer.set(windowExpires);
+      Instant windowEnds = window.maxTimestamp().plus(allowedLateness);
+      LOG.debug("*** SET TIMER *** to point in time {} for window {}", 
windowEnds, window);
+      windowTimer.set(windowEnds);
       key.write(element.getKey());
+      LOG.debug("*** BATCH *** Add element for window {} ", window);
       batch.add(element.getValue());
-      LOG.debug("*** BATCH *** Add element for window {} ", window.toString());
-      // blind add is supported with combiningState
+      // Blind add is supported with combiningState
       numElementsInBatch.add(1L);
+
       Long num = numElementsInBatch.read();
+      if (num == 1 && maxBufferingDuration.isLongerThan(Duration.ZERO)) {
+        // This is the first element in batch. Start counting buffering time 
if a limit was set.
+        bufferingTimer.offset(maxBufferingDuration).setRelative();
+      }
       if (num % prefetchFrequency == 0) {
-        // prefetch data and modify batch state (readLater() modifies this)
+        // Prefetch data and modify batch state (readLater() modifies this)
         batch.readLater();
       }
       if (num >= batchSize) {
         LOG.debug("*** END OF BATCH *** for window {}", window.toString());
-        flushBatch(receiver, key, batch, numElementsInBatch);
+        flushBatch(receiver, key, batch, numElementsInBatch, bufferingTimer);
       }
     }
 
+    @OnTimer(END_OF_BUFFERING_ID)
+    public void onBufferingTimer(
+        OutputReceiver<KV<K, Iterable<InputT>>> receiver,
+        @Timestamp Instant timestamp,
+        @StateId(KEY_ID) ValueState<K> key,
+        @StateId(BATCH_ID) BagState<InputT> batch,
+        @StateId(NUM_ELEMENTS_IN_BATCH_ID) CombiningState<Long, long[], Long> 
numElementsInBatch,
+        @TimerId(END_OF_BUFFERING_ID) Timer bufferingTimer) {
+      LOG.debug(
+          "*** END OF BUFFERING *** for timer timestamp {} with buffering 
duration {}",
+          timestamp,
+          maxBufferingDuration);
+      flushBatch(receiver, key, batch, numElementsInBatch, bufferingTimer);

Review comment:
       We want to pass `bufferingTimer` to null when `flushBatch`is called here 
since there is no need to clear the bufferTimer when it's fired.

##########
File path: 
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupIntoBatches.java
##########
@@ -150,70 +185,96 @@ public long apply(long left, long right) {
               });
 
       this.keySpec = StateSpecs.value(inputKeyCoder);
-      // prefetch every 20% of batchSize elements. Do not prefetch if 
batchSize is too little
+      // Prefetch every 20% of batchSize elements. Do not prefetch if 
batchSize is too little
       this.prefetchFrequency = ((batchSize / 5) <= 1) ? Long.MAX_VALUE : 
(batchSize / 5);
     }
 
     @ProcessElement
     public void processElement(
-        @TimerId(END_OF_WINDOW_ID) Timer timer,
+        @TimerId(END_OF_WINDOW_ID) Timer windowTimer,
+        @TimerId(END_OF_BUFFERING_ID) Timer bufferingTimer,
         @StateId(BATCH_ID) BagState<InputT> batch,
         @StateId(NUM_ELEMENTS_IN_BATCH_ID) CombiningState<Long, long[], Long> 
numElementsInBatch,
         @StateId(KEY_ID) ValueState<K> key,
         @Element KV<K, InputT> element,
         BoundedWindow window,
         OutputReceiver<KV<K, Iterable<InputT>>> receiver) {
-      Instant windowExpires = window.maxTimestamp().plus(allowedLateness);
-
-      LOG.debug(
-          "*** SET TIMER *** to point in time {} for window {}",
-          windowExpires.toString(),
-          window.toString());
-      timer.set(windowExpires);
+      Instant windowEnds = window.maxTimestamp().plus(allowedLateness);
+      LOG.debug("*** SET TIMER *** to point in time {} for window {}", 
windowEnds, window);
+      windowTimer.set(windowEnds);
       key.write(element.getKey());
+      LOG.debug("*** BATCH *** Add element for window {} ", window);
       batch.add(element.getValue());
-      LOG.debug("*** BATCH *** Add element for window {} ", window.toString());
-      // blind add is supported with combiningState
+      // Blind add is supported with combiningState
       numElementsInBatch.add(1L);
+
       Long num = numElementsInBatch.read();
+      if (num == 1 && maxBufferingDuration.isLongerThan(Duration.ZERO)) {

Review comment:
       If we `checkArgument(duration != Duration.ZERO)` when 
`withMaxBufferingDuration ()`, then we can drop the check here.

##########
File path: 
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupIntoBatches.java
##########
@@ -66,25 +68,49 @@
  *         }
  *     }}));
  * </pre>
+ *
+ * <p>Example 2 (batch unbounded input in a global window):
+ *
+ * <pre>{@code
+ * PCollection<KV<String, String>> unboundedInput = ...;
+ * long batchSize = 100L;
+ * Duration maxBufferingDuration = Duration.standardSeconds(10);
+ * PCollection<KV<String, Iterable<String>>> batched = unboundedInput
+ *     .apply(Window.<KV<String, String>>into(new GlobalWindows())
+ *         .triggering(Repeatedly.forever(AfterPane.elementCountAtLeast(1)))
+ *         .discardingFiredPanes())
+ *     .apply(GroupIntoBatches.<String, String>ofSize(batchSize)
+ *         .withMaxBufferingDuration(maxBufferingDuration));
+ * }</pre>
  */
 public class GroupIntoBatches<K, InputT>
     extends PTransform<PCollection<KV<K, InputT>>, PCollection<KV<K, 
Iterable<InputT>>>> {
 
   private final long batchSize;
+  private final Duration maxBufferingDuration;
 
-  private GroupIntoBatches(long batchSize) {
+  private GroupIntoBatches(long batchSize, Duration maxBufferingDuration) {
     this.batchSize = batchSize;
+    this.maxBufferingDuration = maxBufferingDuration;
   }
 
   public static <K, InputT> GroupIntoBatches<K, InputT> ofSize(long batchSize) 
{
-    return new GroupIntoBatches<>(batchSize);
+    return new GroupIntoBatches<>(batchSize, Duration.ZERO);
   }
 
   /** Returns the size of the batch. */
   public long getBatchSize() {
     return batchSize;
   }
 
+  /**
+   * Set a time limit (in processing time) on how long an incomplete batch of 
elements is allowed to
+   * be buffered. Once a batch is flushed to output, the timer is reset.
+   */
+  public GroupIntoBatches<K, InputT> withMaxBufferingDuration(Duration 
duration) {
+    return new GroupIntoBatches<>(batchSize, duration);

Review comment:
       If `Duration.ZERO` is not preferred, we can have a check here.

##########
File path: 
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupIntoBatches.java
##########
@@ -150,70 +185,96 @@ public long apply(long left, long right) {
               });
 
       this.keySpec = StateSpecs.value(inputKeyCoder);
-      // prefetch every 20% of batchSize elements. Do not prefetch if 
batchSize is too little
+      // Prefetch every 20% of batchSize elements. Do not prefetch if 
batchSize is too little
       this.prefetchFrequency = ((batchSize / 5) <= 1) ? Long.MAX_VALUE : 
(batchSize / 5);
     }
 
     @ProcessElement
     public void processElement(
-        @TimerId(END_OF_WINDOW_ID) Timer timer,
+        @TimerId(END_OF_WINDOW_ID) Timer windowTimer,
+        @TimerId(END_OF_BUFFERING_ID) Timer bufferingTimer,
         @StateId(BATCH_ID) BagState<InputT> batch,
         @StateId(NUM_ELEMENTS_IN_BATCH_ID) CombiningState<Long, long[], Long> 
numElementsInBatch,
         @StateId(KEY_ID) ValueState<K> key,
         @Element KV<K, InputT> element,
         BoundedWindow window,
         OutputReceiver<KV<K, Iterable<InputT>>> receiver) {
-      Instant windowExpires = window.maxTimestamp().plus(allowedLateness);
-
-      LOG.debug(
-          "*** SET TIMER *** to point in time {} for window {}",
-          windowExpires.toString(),
-          window.toString());
-      timer.set(windowExpires);
+      Instant windowEnds = window.maxTimestamp().plus(allowedLateness);
+      LOG.debug("*** SET TIMER *** to point in time {} for window {}", 
windowEnds, window);
+      windowTimer.set(windowEnds);
       key.write(element.getKey());
+      LOG.debug("*** BATCH *** Add element for window {} ", window);
       batch.add(element.getValue());
-      LOG.debug("*** BATCH *** Add element for window {} ", window.toString());
-      // blind add is supported with combiningState
+      // Blind add is supported with combiningState
       numElementsInBatch.add(1L);
+
       Long num = numElementsInBatch.read();
+      if (num == 1 && maxBufferingDuration.isLongerThan(Duration.ZERO)) {
+        // This is the first element in batch. Start counting buffering time 
if a limit was set.
+        bufferingTimer.offset(maxBufferingDuration).setRelative();
+      }
       if (num % prefetchFrequency == 0) {
-        // prefetch data and modify batch state (readLater() modifies this)
+        // Prefetch data and modify batch state (readLater() modifies this)
         batch.readLater();
       }
       if (num >= batchSize) {
         LOG.debug("*** END OF BATCH *** for window {}", window.toString());
-        flushBatch(receiver, key, batch, numElementsInBatch);
+        flushBatch(receiver, key, batch, numElementsInBatch, bufferingTimer);
       }
     }
 
+    @OnTimer(END_OF_BUFFERING_ID)
+    public void onBufferingTimer(
+        OutputReceiver<KV<K, Iterable<InputT>>> receiver,
+        @Timestamp Instant timestamp,
+        @StateId(KEY_ID) ValueState<K> key,
+        @StateId(BATCH_ID) BagState<InputT> batch,
+        @StateId(NUM_ELEMENTS_IN_BATCH_ID) CombiningState<Long, long[], Long> 
numElementsInBatch,
+        @TimerId(END_OF_BUFFERING_ID) Timer bufferingTimer) {
+      LOG.debug(
+          "*** END OF BUFFERING *** for timer timestamp {} with buffering 
duration {}",
+          timestamp,
+          maxBufferingDuration);
+      flushBatch(receiver, key, batch, numElementsInBatch, bufferingTimer);
+    }
+
     @OnTimer(END_OF_WINDOW_ID)
-    public void onTimerCallback(
+    public void onWindowTimer(
         OutputReceiver<KV<K, Iterable<InputT>>> receiver,
         @Timestamp Instant timestamp,
         @StateId(KEY_ID) ValueState<K> key,
         @StateId(BATCH_ID) BagState<InputT> batch,
         @StateId(NUM_ELEMENTS_IN_BATCH_ID) CombiningState<Long, long[], Long> 
numElementsInBatch,
+        @TimerId(END_OF_BUFFERING_ID) Timer bufferingTimer,
         BoundedWindow window) {
       LOG.debug(
           "*** END OF WINDOW *** for timer timestamp {} in windows {}",
           timestamp,
           window.toString());
-      flushBatch(receiver, key, batch, numElementsInBatch);
+      flushBatch(receiver, key, batch, numElementsInBatch, null);
     }
 
     private void flushBatch(
         OutputReceiver<KV<K, Iterable<InputT>>> receiver,
         ValueState<K> key,
         BagState<InputT> batch,
-        CombiningState<Long, long[], Long> numElementsInBatch) {
+        CombiningState<Long, long[], Long> numElementsInBatch,
+        @Nullable Timer bufferingTimer) {
       Iterable<InputT> values = batch.read();
-      // when the timer fires, batch state might be empty
+      // When the timer fires, batch state might be empty
       if (!Iterables.isEmpty(values)) {
         receiver.output(KV.of(key.read(), values));
       }
       batch.clear();
       LOG.debug("*** BATCH *** clear");
       numElementsInBatch.clear();
+      // We might reach here due to batch size being reached or window 
expiration. Reset the
+      // buffering timer (if not null) since the state is empty now. It'll be 
extended again if a
+      // new element arrives prior to the expiration time set here.
+      // TODO(BEAM-10887): Use clear() when it's available.
+      if (bufferingTimer != null && 
maxBufferingDuration.isLongerThan(Duration.ZERO)) {

Review comment:
       Same above.

##########
File path: 
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupIntoBatches.java
##########
@@ -150,70 +185,96 @@ public long apply(long left, long right) {
               });
 
       this.keySpec = StateSpecs.value(inputKeyCoder);
-      // prefetch every 20% of batchSize elements. Do not prefetch if 
batchSize is too little
+      // Prefetch every 20% of batchSize elements. Do not prefetch if 
batchSize is too little
       this.prefetchFrequency = ((batchSize / 5) <= 1) ? Long.MAX_VALUE : 
(batchSize / 5);
     }
 
     @ProcessElement
     public void processElement(
-        @TimerId(END_OF_WINDOW_ID) Timer timer,
+        @TimerId(END_OF_WINDOW_ID) Timer windowTimer,
+        @TimerId(END_OF_BUFFERING_ID) Timer bufferingTimer,
         @StateId(BATCH_ID) BagState<InputT> batch,
         @StateId(NUM_ELEMENTS_IN_BATCH_ID) CombiningState<Long, long[], Long> 
numElementsInBatch,
         @StateId(KEY_ID) ValueState<K> key,
         @Element KV<K, InputT> element,
         BoundedWindow window,
         OutputReceiver<KV<K, Iterable<InputT>>> receiver) {
-      Instant windowExpires = window.maxTimestamp().plus(allowedLateness);
-
-      LOG.debug(
-          "*** SET TIMER *** to point in time {} for window {}",
-          windowExpires.toString(),
-          window.toString());
-      timer.set(windowExpires);
+      Instant windowEnds = window.maxTimestamp().plus(allowedLateness);
+      LOG.debug("*** SET TIMER *** to point in time {} for window {}", 
windowEnds, window);
+      windowTimer.set(windowEnds);
       key.write(element.getKey());
+      LOG.debug("*** BATCH *** Add element for window {} ", window);
       batch.add(element.getValue());
-      LOG.debug("*** BATCH *** Add element for window {} ", window.toString());
-      // blind add is supported with combiningState
+      // Blind add is supported with combiningState
       numElementsInBatch.add(1L);
+
       Long num = numElementsInBatch.read();
+      if (num == 1 && maxBufferingDuration.isLongerThan(Duration.ZERO)) {
+        // This is the first element in batch. Start counting buffering time 
if a limit was set.
+        bufferingTimer.offset(maxBufferingDuration).setRelative();
+      }
       if (num % prefetchFrequency == 0) {
-        // prefetch data and modify batch state (readLater() modifies this)
+        // Prefetch data and modify batch state (readLater() modifies this)
         batch.readLater();
       }
       if (num >= batchSize) {
         LOG.debug("*** END OF BATCH *** for window {}", window.toString());
-        flushBatch(receiver, key, batch, numElementsInBatch);
+        flushBatch(receiver, key, batch, numElementsInBatch, bufferingTimer);
       }
     }
 
+    @OnTimer(END_OF_BUFFERING_ID)
+    public void onBufferingTimer(
+        OutputReceiver<KV<K, Iterable<InputT>>> receiver,
+        @Timestamp Instant timestamp,
+        @StateId(KEY_ID) ValueState<K> key,
+        @StateId(BATCH_ID) BagState<InputT> batch,
+        @StateId(NUM_ELEMENTS_IN_BATCH_ID) CombiningState<Long, long[], Long> 
numElementsInBatch,
+        @TimerId(END_OF_BUFFERING_ID) Timer bufferingTimer) {
+      LOG.debug(
+          "*** END OF BUFFERING *** for timer timestamp {} with buffering 
duration {}",
+          timestamp,
+          maxBufferingDuration);
+      flushBatch(receiver, key, batch, numElementsInBatch, bufferingTimer);
+    }
+
     @OnTimer(END_OF_WINDOW_ID)
-    public void onTimerCallback(
+    public void onWindowTimer(
         OutputReceiver<KV<K, Iterable<InputT>>> receiver,
         @Timestamp Instant timestamp,
         @StateId(KEY_ID) ValueState<K> key,
         @StateId(BATCH_ID) BagState<InputT> batch,
         @StateId(NUM_ELEMENTS_IN_BATCH_ID) CombiningState<Long, long[], Long> 
numElementsInBatch,
+        @TimerId(END_OF_BUFFERING_ID) Timer bufferingTimer,
         BoundedWindow window) {
       LOG.debug(
           "*** END OF WINDOW *** for timer timestamp {} in windows {}",
           timestamp,
           window.toString());
-      flushBatch(receiver, key, batch, numElementsInBatch);
+      flushBatch(receiver, key, batch, numElementsInBatch, null);

Review comment:
       Instead, we want to pass actual bufferingTimer here since we want to 
clear it.

##########
File path: 
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupIntoBatches.java
##########
@@ -66,25 +68,49 @@
  *         }
  *     }}));
  * </pre>
+ *
+ * <p>Example 2 (batch unbounded input in a global window):
+ *
+ * <pre>{@code
+ * PCollection<KV<String, String>> unboundedInput = ...;
+ * long batchSize = 100L;
+ * Duration maxBufferingDuration = Duration.standardSeconds(10);
+ * PCollection<KV<String, Iterable<String>>> batched = unboundedInput
+ *     .apply(Window.<KV<String, String>>into(new GlobalWindows())
+ *         .triggering(Repeatedly.forever(AfterPane.elementCountAtLeast(1)))
+ *         .discardingFiredPanes())
+ *     .apply(GroupIntoBatches.<String, String>ofSize(batchSize)
+ *         .withMaxBufferingDuration(maxBufferingDuration));
+ * }</pre>
  */
 public class GroupIntoBatches<K, InputT>
     extends PTransform<PCollection<KV<K, InputT>>, PCollection<KV<K, 
Iterable<InputT>>>> {
 
   private final long batchSize;
+  private final Duration maxBufferingDuration;
 
-  private GroupIntoBatches(long batchSize) {
+  private GroupIntoBatches(long batchSize, Duration maxBufferingDuration) {
     this.batchSize = batchSize;
+    this.maxBufferingDuration = maxBufferingDuration;
   }
 
   public static <K, InputT> GroupIntoBatches<K, InputT> ofSize(long batchSize) 
{
-    return new GroupIntoBatches<>(batchSize);
+    return new GroupIntoBatches<>(batchSize, Duration.ZERO);

Review comment:
       Can we use `null` here?

##########
File path: 
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupIntoBatches.java
##########
@@ -150,70 +185,96 @@ public long apply(long left, long right) {
               });
 
       this.keySpec = StateSpecs.value(inputKeyCoder);
-      // prefetch every 20% of batchSize elements. Do not prefetch if 
batchSize is too little
+      // Prefetch every 20% of batchSize elements. Do not prefetch if 
batchSize is too little
       this.prefetchFrequency = ((batchSize / 5) <= 1) ? Long.MAX_VALUE : 
(batchSize / 5);
     }
 
     @ProcessElement
     public void processElement(
-        @TimerId(END_OF_WINDOW_ID) Timer timer,
+        @TimerId(END_OF_WINDOW_ID) Timer windowTimer,
+        @TimerId(END_OF_BUFFERING_ID) Timer bufferingTimer,
         @StateId(BATCH_ID) BagState<InputT> batch,
         @StateId(NUM_ELEMENTS_IN_BATCH_ID) CombiningState<Long, long[], Long> 
numElementsInBatch,
         @StateId(KEY_ID) ValueState<K> key,
         @Element KV<K, InputT> element,
         BoundedWindow window,
         OutputReceiver<KV<K, Iterable<InputT>>> receiver) {
-      Instant windowExpires = window.maxTimestamp().plus(allowedLateness);
-
-      LOG.debug(
-          "*** SET TIMER *** to point in time {} for window {}",
-          windowExpires.toString(),
-          window.toString());
-      timer.set(windowExpires);
+      Instant windowEnds = window.maxTimestamp().plus(allowedLateness);
+      LOG.debug("*** SET TIMER *** to point in time {} for window {}", 
windowEnds, window);
+      windowTimer.set(windowEnds);
       key.write(element.getKey());
+      LOG.debug("*** BATCH *** Add element for window {} ", window);
       batch.add(element.getValue());
-      LOG.debug("*** BATCH *** Add element for window {} ", window.toString());
-      // blind add is supported with combiningState
+      // Blind add is supported with combiningState
       numElementsInBatch.add(1L);
+
       Long num = numElementsInBatch.read();
+      if (num == 1 && maxBufferingDuration.isLongerThan(Duration.ZERO)) {
+        // This is the first element in batch. Start counting buffering time 
if a limit was set.
+        bufferingTimer.offset(maxBufferingDuration).setRelative();
+      }
       if (num % prefetchFrequency == 0) {
-        // prefetch data and modify batch state (readLater() modifies this)
+        // Prefetch data and modify batch state (readLater() modifies this)
         batch.readLater();
       }
       if (num >= batchSize) {
         LOG.debug("*** END OF BATCH *** for window {}", window.toString());
-        flushBatch(receiver, key, batch, numElementsInBatch);
+        flushBatch(receiver, key, batch, numElementsInBatch, bufferingTimer);
       }
     }
 
+    @OnTimer(END_OF_BUFFERING_ID)
+    public void onBufferingTimer(
+        OutputReceiver<KV<K, Iterable<InputT>>> receiver,
+        @Timestamp Instant timestamp,
+        @StateId(KEY_ID) ValueState<K> key,
+        @StateId(BATCH_ID) BagState<InputT> batch,
+        @StateId(NUM_ELEMENTS_IN_BATCH_ID) CombiningState<Long, long[], Long> 
numElementsInBatch,
+        @TimerId(END_OF_BUFFERING_ID) Timer bufferingTimer) {
+      LOG.debug(
+          "*** END OF BUFFERING *** for timer timestamp {} with buffering 
duration {}",
+          timestamp,
+          maxBufferingDuration);
+      flushBatch(receiver, key, batch, numElementsInBatch, bufferingTimer);

Review comment:
       We want to pass `bufferingTimer` to null when `flushBatch`is called here 
since there is no need to clear the bufferTimer when it's fired.

##########
File path: 
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupIntoBatches.java
##########
@@ -150,70 +185,96 @@ public long apply(long left, long right) {
               });
 
       this.keySpec = StateSpecs.value(inputKeyCoder);
-      // prefetch every 20% of batchSize elements. Do not prefetch if 
batchSize is too little
+      // Prefetch every 20% of batchSize elements. Do not prefetch if 
batchSize is too little
       this.prefetchFrequency = ((batchSize / 5) <= 1) ? Long.MAX_VALUE : 
(batchSize / 5);
     }
 
     @ProcessElement
     public void processElement(
-        @TimerId(END_OF_WINDOW_ID) Timer timer,
+        @TimerId(END_OF_WINDOW_ID) Timer windowTimer,
+        @TimerId(END_OF_BUFFERING_ID) Timer bufferingTimer,
         @StateId(BATCH_ID) BagState<InputT> batch,
         @StateId(NUM_ELEMENTS_IN_BATCH_ID) CombiningState<Long, long[], Long> 
numElementsInBatch,
         @StateId(KEY_ID) ValueState<K> key,
         @Element KV<K, InputT> element,
         BoundedWindow window,
         OutputReceiver<KV<K, Iterable<InputT>>> receiver) {
-      Instant windowExpires = window.maxTimestamp().plus(allowedLateness);
-
-      LOG.debug(
-          "*** SET TIMER *** to point in time {} for window {}",
-          windowExpires.toString(),
-          window.toString());
-      timer.set(windowExpires);
+      Instant windowEnds = window.maxTimestamp().plus(allowedLateness);
+      LOG.debug("*** SET TIMER *** to point in time {} for window {}", 
windowEnds, window);
+      windowTimer.set(windowEnds);
       key.write(element.getKey());
+      LOG.debug("*** BATCH *** Add element for window {} ", window);
       batch.add(element.getValue());
-      LOG.debug("*** BATCH *** Add element for window {} ", window.toString());
-      // blind add is supported with combiningState
+      // Blind add is supported with combiningState
       numElementsInBatch.add(1L);
+
       Long num = numElementsInBatch.read();
+      if (num == 1 && maxBufferingDuration.isLongerThan(Duration.ZERO)) {

Review comment:
       If we `checkArgument(duration != Duration.ZERO)` when 
`withMaxBufferingDuration ()`, then we can drop the check here.

##########
File path: 
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupIntoBatches.java
##########
@@ -66,25 +68,49 @@
  *         }
  *     }}));
  * </pre>
+ *
+ * <p>Example 2 (batch unbounded input in a global window):
+ *
+ * <pre>{@code
+ * PCollection<KV<String, String>> unboundedInput = ...;
+ * long batchSize = 100L;
+ * Duration maxBufferingDuration = Duration.standardSeconds(10);
+ * PCollection<KV<String, Iterable<String>>> batched = unboundedInput
+ *     .apply(Window.<KV<String, String>>into(new GlobalWindows())
+ *         .triggering(Repeatedly.forever(AfterPane.elementCountAtLeast(1)))
+ *         .discardingFiredPanes())
+ *     .apply(GroupIntoBatches.<String, String>ofSize(batchSize)
+ *         .withMaxBufferingDuration(maxBufferingDuration));
+ * }</pre>
  */
 public class GroupIntoBatches<K, InputT>
     extends PTransform<PCollection<KV<K, InputT>>, PCollection<KV<K, 
Iterable<InputT>>>> {
 
   private final long batchSize;
+  private final Duration maxBufferingDuration;
 
-  private GroupIntoBatches(long batchSize) {
+  private GroupIntoBatches(long batchSize, Duration maxBufferingDuration) {
     this.batchSize = batchSize;
+    this.maxBufferingDuration = maxBufferingDuration;
   }
 
   public static <K, InputT> GroupIntoBatches<K, InputT> ofSize(long batchSize) 
{
-    return new GroupIntoBatches<>(batchSize);
+    return new GroupIntoBatches<>(batchSize, Duration.ZERO);
   }
 
   /** Returns the size of the batch. */
   public long getBatchSize() {
     return batchSize;
   }
 
+  /**
+   * Set a time limit (in processing time) on how long an incomplete batch of 
elements is allowed to
+   * be buffered. Once a batch is flushed to output, the timer is reset.
+   */
+  public GroupIntoBatches<K, InputT> withMaxBufferingDuration(Duration 
duration) {
+    return new GroupIntoBatches<>(batchSize, duration);

Review comment:
       If `Duration.ZERO` is not preferred, we can have a check here.

##########
File path: 
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupIntoBatches.java
##########
@@ -150,70 +185,96 @@ public long apply(long left, long right) {
               });
 
       this.keySpec = StateSpecs.value(inputKeyCoder);
-      // prefetch every 20% of batchSize elements. Do not prefetch if 
batchSize is too little
+      // Prefetch every 20% of batchSize elements. Do not prefetch if 
batchSize is too little
       this.prefetchFrequency = ((batchSize / 5) <= 1) ? Long.MAX_VALUE : 
(batchSize / 5);
     }
 
     @ProcessElement
     public void processElement(
-        @TimerId(END_OF_WINDOW_ID) Timer timer,
+        @TimerId(END_OF_WINDOW_ID) Timer windowTimer,
+        @TimerId(END_OF_BUFFERING_ID) Timer bufferingTimer,
         @StateId(BATCH_ID) BagState<InputT> batch,
         @StateId(NUM_ELEMENTS_IN_BATCH_ID) CombiningState<Long, long[], Long> 
numElementsInBatch,
         @StateId(KEY_ID) ValueState<K> key,
         @Element KV<K, InputT> element,
         BoundedWindow window,
         OutputReceiver<KV<K, Iterable<InputT>>> receiver) {
-      Instant windowExpires = window.maxTimestamp().plus(allowedLateness);
-
-      LOG.debug(
-          "*** SET TIMER *** to point in time {} for window {}",
-          windowExpires.toString(),
-          window.toString());
-      timer.set(windowExpires);
+      Instant windowEnds = window.maxTimestamp().plus(allowedLateness);
+      LOG.debug("*** SET TIMER *** to point in time {} for window {}", 
windowEnds, window);
+      windowTimer.set(windowEnds);
       key.write(element.getKey());
+      LOG.debug("*** BATCH *** Add element for window {} ", window);
       batch.add(element.getValue());
-      LOG.debug("*** BATCH *** Add element for window {} ", window.toString());
-      // blind add is supported with combiningState
+      // Blind add is supported with combiningState
       numElementsInBatch.add(1L);
+
       Long num = numElementsInBatch.read();
+      if (num == 1 && maxBufferingDuration.isLongerThan(Duration.ZERO)) {
+        // This is the first element in batch. Start counting buffering time 
if a limit was set.
+        bufferingTimer.offset(maxBufferingDuration).setRelative();
+      }
       if (num % prefetchFrequency == 0) {
-        // prefetch data and modify batch state (readLater() modifies this)
+        // Prefetch data and modify batch state (readLater() modifies this)
         batch.readLater();
       }
       if (num >= batchSize) {
         LOG.debug("*** END OF BATCH *** for window {}", window.toString());
-        flushBatch(receiver, key, batch, numElementsInBatch);
+        flushBatch(receiver, key, batch, numElementsInBatch, bufferingTimer);
       }
     }
 
+    @OnTimer(END_OF_BUFFERING_ID)
+    public void onBufferingTimer(
+        OutputReceiver<KV<K, Iterable<InputT>>> receiver,
+        @Timestamp Instant timestamp,
+        @StateId(KEY_ID) ValueState<K> key,
+        @StateId(BATCH_ID) BagState<InputT> batch,
+        @StateId(NUM_ELEMENTS_IN_BATCH_ID) CombiningState<Long, long[], Long> 
numElementsInBatch,
+        @TimerId(END_OF_BUFFERING_ID) Timer bufferingTimer) {
+      LOG.debug(
+          "*** END OF BUFFERING *** for timer timestamp {} with buffering 
duration {}",
+          timestamp,
+          maxBufferingDuration);
+      flushBatch(receiver, key, batch, numElementsInBatch, bufferingTimer);
+    }
+
     @OnTimer(END_OF_WINDOW_ID)
-    public void onTimerCallback(
+    public void onWindowTimer(
         OutputReceiver<KV<K, Iterable<InputT>>> receiver,
         @Timestamp Instant timestamp,
         @StateId(KEY_ID) ValueState<K> key,
         @StateId(BATCH_ID) BagState<InputT> batch,
         @StateId(NUM_ELEMENTS_IN_BATCH_ID) CombiningState<Long, long[], Long> 
numElementsInBatch,
+        @TimerId(END_OF_BUFFERING_ID) Timer bufferingTimer,
         BoundedWindow window) {
       LOG.debug(
           "*** END OF WINDOW *** for timer timestamp {} in windows {}",
           timestamp,
           window.toString());
-      flushBatch(receiver, key, batch, numElementsInBatch);
+      flushBatch(receiver, key, batch, numElementsInBatch, null);
     }
 
     private void flushBatch(
         OutputReceiver<KV<K, Iterable<InputT>>> receiver,
         ValueState<K> key,
         BagState<InputT> batch,
-        CombiningState<Long, long[], Long> numElementsInBatch) {
+        CombiningState<Long, long[], Long> numElementsInBatch,
+        @Nullable Timer bufferingTimer) {
       Iterable<InputT> values = batch.read();
-      // when the timer fires, batch state might be empty
+      // When the timer fires, batch state might be empty
       if (!Iterables.isEmpty(values)) {
         receiver.output(KV.of(key.read(), values));
       }
       batch.clear();
       LOG.debug("*** BATCH *** clear");
       numElementsInBatch.clear();
+      // We might reach here due to batch size being reached or window 
expiration. Reset the
+      // buffering timer (if not null) since the state is empty now. It'll be 
extended again if a
+      // new element arrives prior to the expiration time set here.
+      // TODO(BEAM-10887): Use clear() when it's available.
+      if (bufferingTimer != null && 
maxBufferingDuration.isLongerThan(Duration.ZERO)) {

Review comment:
       Same above.

##########
File path: 
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupIntoBatches.java
##########
@@ -150,70 +185,96 @@ public long apply(long left, long right) {
               });
 
       this.keySpec = StateSpecs.value(inputKeyCoder);
-      // prefetch every 20% of batchSize elements. Do not prefetch if 
batchSize is too little
+      // Prefetch every 20% of batchSize elements. Do not prefetch if 
batchSize is too little
       this.prefetchFrequency = ((batchSize / 5) <= 1) ? Long.MAX_VALUE : 
(batchSize / 5);
     }
 
     @ProcessElement
     public void processElement(
-        @TimerId(END_OF_WINDOW_ID) Timer timer,
+        @TimerId(END_OF_WINDOW_ID) Timer windowTimer,
+        @TimerId(END_OF_BUFFERING_ID) Timer bufferingTimer,
         @StateId(BATCH_ID) BagState<InputT> batch,
         @StateId(NUM_ELEMENTS_IN_BATCH_ID) CombiningState<Long, long[], Long> 
numElementsInBatch,
         @StateId(KEY_ID) ValueState<K> key,
         @Element KV<K, InputT> element,
         BoundedWindow window,
         OutputReceiver<KV<K, Iterable<InputT>>> receiver) {
-      Instant windowExpires = window.maxTimestamp().plus(allowedLateness);
-
-      LOG.debug(
-          "*** SET TIMER *** to point in time {} for window {}",
-          windowExpires.toString(),
-          window.toString());
-      timer.set(windowExpires);
+      Instant windowEnds = window.maxTimestamp().plus(allowedLateness);
+      LOG.debug("*** SET TIMER *** to point in time {} for window {}", 
windowEnds, window);
+      windowTimer.set(windowEnds);
       key.write(element.getKey());
+      LOG.debug("*** BATCH *** Add element for window {} ", window);
       batch.add(element.getValue());
-      LOG.debug("*** BATCH *** Add element for window {} ", window.toString());
-      // blind add is supported with combiningState
+      // Blind add is supported with combiningState
       numElementsInBatch.add(1L);
+
       Long num = numElementsInBatch.read();
+      if (num == 1 && maxBufferingDuration.isLongerThan(Duration.ZERO)) {
+        // This is the first element in batch. Start counting buffering time 
if a limit was set.
+        bufferingTimer.offset(maxBufferingDuration).setRelative();
+      }
       if (num % prefetchFrequency == 0) {
-        // prefetch data and modify batch state (readLater() modifies this)
+        // Prefetch data and modify batch state (readLater() modifies this)
         batch.readLater();
       }
       if (num >= batchSize) {
         LOG.debug("*** END OF BATCH *** for window {}", window.toString());
-        flushBatch(receiver, key, batch, numElementsInBatch);
+        flushBatch(receiver, key, batch, numElementsInBatch, bufferingTimer);
       }
     }
 
+    @OnTimer(END_OF_BUFFERING_ID)
+    public void onBufferingTimer(
+        OutputReceiver<KV<K, Iterable<InputT>>> receiver,
+        @Timestamp Instant timestamp,
+        @StateId(KEY_ID) ValueState<K> key,
+        @StateId(BATCH_ID) BagState<InputT> batch,
+        @StateId(NUM_ELEMENTS_IN_BATCH_ID) CombiningState<Long, long[], Long> 
numElementsInBatch,
+        @TimerId(END_OF_BUFFERING_ID) Timer bufferingTimer) {
+      LOG.debug(
+          "*** END OF BUFFERING *** for timer timestamp {} with buffering 
duration {}",
+          timestamp,
+          maxBufferingDuration);
+      flushBatch(receiver, key, batch, numElementsInBatch, bufferingTimer);
+    }
+
     @OnTimer(END_OF_WINDOW_ID)
-    public void onTimerCallback(
+    public void onWindowTimer(
         OutputReceiver<KV<K, Iterable<InputT>>> receiver,
         @Timestamp Instant timestamp,
         @StateId(KEY_ID) ValueState<K> key,
         @StateId(BATCH_ID) BagState<InputT> batch,
         @StateId(NUM_ELEMENTS_IN_BATCH_ID) CombiningState<Long, long[], Long> 
numElementsInBatch,
+        @TimerId(END_OF_BUFFERING_ID) Timer bufferingTimer,
         BoundedWindow window) {
       LOG.debug(
           "*** END OF WINDOW *** for timer timestamp {} in windows {}",
           timestamp,
           window.toString());
-      flushBatch(receiver, key, batch, numElementsInBatch);
+      flushBatch(receiver, key, batch, numElementsInBatch, null);

Review comment:
       Instead, we want to pass actual bufferingTimer here since we want to 
clear it.

##########
File path: 
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupIntoBatches.java
##########
@@ -66,25 +68,49 @@
  *         }
  *     }}));
  * </pre>
+ *
+ * <p>Example 2 (batch unbounded input in a global window):
+ *
+ * <pre>{@code
+ * PCollection<KV<String, String>> unboundedInput = ...;
+ * long batchSize = 100L;
+ * Duration maxBufferingDuration = Duration.standardSeconds(10);
+ * PCollection<KV<String, Iterable<String>>> batched = unboundedInput
+ *     .apply(Window.<KV<String, String>>into(new GlobalWindows())
+ *         .triggering(Repeatedly.forever(AfterPane.elementCountAtLeast(1)))
+ *         .discardingFiredPanes())
+ *     .apply(GroupIntoBatches.<String, String>ofSize(batchSize)
+ *         .withMaxBufferingDuration(maxBufferingDuration));
+ * }</pre>
  */
 public class GroupIntoBatches<K, InputT>
     extends PTransform<PCollection<KV<K, InputT>>, PCollection<KV<K, 
Iterable<InputT>>>> {
 
   private final long batchSize;
+  private final Duration maxBufferingDuration;
 
-  private GroupIntoBatches(long batchSize) {
+  private GroupIntoBatches(long batchSize, Duration maxBufferingDuration) {
     this.batchSize = batchSize;
+    this.maxBufferingDuration = maxBufferingDuration;
   }
 
   public static <K, InputT> GroupIntoBatches<K, InputT> ofSize(long batchSize) 
{
-    return new GroupIntoBatches<>(batchSize);
+    return new GroupIntoBatches<>(batchSize, Duration.ZERO);

Review comment:
       Can we use `null` here?

##########
File path: 
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupIntoBatches.java
##########
@@ -150,70 +185,96 @@ public long apply(long left, long right) {
               });
 
       this.keySpec = StateSpecs.value(inputKeyCoder);
-      // prefetch every 20% of batchSize elements. Do not prefetch if 
batchSize is too little
+      // Prefetch every 20% of batchSize elements. Do not prefetch if 
batchSize is too little
       this.prefetchFrequency = ((batchSize / 5) <= 1) ? Long.MAX_VALUE : 
(batchSize / 5);
     }
 
     @ProcessElement
     public void processElement(
-        @TimerId(END_OF_WINDOW_ID) Timer timer,
+        @TimerId(END_OF_WINDOW_ID) Timer windowTimer,
+        @TimerId(END_OF_BUFFERING_ID) Timer bufferingTimer,
         @StateId(BATCH_ID) BagState<InputT> batch,
         @StateId(NUM_ELEMENTS_IN_BATCH_ID) CombiningState<Long, long[], Long> 
numElementsInBatch,
         @StateId(KEY_ID) ValueState<K> key,
         @Element KV<K, InputT> element,
         BoundedWindow window,
         OutputReceiver<KV<K, Iterable<InputT>>> receiver) {
-      Instant windowExpires = window.maxTimestamp().plus(allowedLateness);
-
-      LOG.debug(
-          "*** SET TIMER *** to point in time {} for window {}",
-          windowExpires.toString(),
-          window.toString());
-      timer.set(windowExpires);
+      Instant windowEnds = window.maxTimestamp().plus(allowedLateness);
+      LOG.debug("*** SET TIMER *** to point in time {} for window {}", 
windowEnds, window);
+      windowTimer.set(windowEnds);
       key.write(element.getKey());
+      LOG.debug("*** BATCH *** Add element for window {} ", window);
       batch.add(element.getValue());
-      LOG.debug("*** BATCH *** Add element for window {} ", window.toString());
-      // blind add is supported with combiningState
+      // Blind add is supported with combiningState
       numElementsInBatch.add(1L);
+
       Long num = numElementsInBatch.read();
+      if (num == 1 && maxBufferingDuration.isLongerThan(Duration.ZERO)) {
+        // This is the first element in batch. Start counting buffering time 
if a limit was set.
+        bufferingTimer.offset(maxBufferingDuration).setRelative();
+      }
       if (num % prefetchFrequency == 0) {
-        // prefetch data and modify batch state (readLater() modifies this)
+        // Prefetch data and modify batch state (readLater() modifies this)
         batch.readLater();
       }
       if (num >= batchSize) {
         LOG.debug("*** END OF BATCH *** for window {}", window.toString());
-        flushBatch(receiver, key, batch, numElementsInBatch);
+        flushBatch(receiver, key, batch, numElementsInBatch, bufferingTimer);
       }
     }
 
+    @OnTimer(END_OF_BUFFERING_ID)
+    public void onBufferingTimer(
+        OutputReceiver<KV<K, Iterable<InputT>>> receiver,
+        @Timestamp Instant timestamp,
+        @StateId(KEY_ID) ValueState<K> key,
+        @StateId(BATCH_ID) BagState<InputT> batch,
+        @StateId(NUM_ELEMENTS_IN_BATCH_ID) CombiningState<Long, long[], Long> 
numElementsInBatch,
+        @TimerId(END_OF_BUFFERING_ID) Timer bufferingTimer) {
+      LOG.debug(
+          "*** END OF BUFFERING *** for timer timestamp {} with buffering 
duration {}",
+          timestamp,
+          maxBufferingDuration);
+      flushBatch(receiver, key, batch, numElementsInBatch, bufferingTimer);

Review comment:
       We want to pass `bufferingTimer` to null when `flushBatch`is called here 
since there is no need to clear the bufferTimer when it's fired.

##########
File path: 
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupIntoBatches.java
##########
@@ -150,70 +185,96 @@ public long apply(long left, long right) {
               });
 
       this.keySpec = StateSpecs.value(inputKeyCoder);
-      // prefetch every 20% of batchSize elements. Do not prefetch if 
batchSize is too little
+      // Prefetch every 20% of batchSize elements. Do not prefetch if 
batchSize is too little
       this.prefetchFrequency = ((batchSize / 5) <= 1) ? Long.MAX_VALUE : 
(batchSize / 5);
     }
 
     @ProcessElement
     public void processElement(
-        @TimerId(END_OF_WINDOW_ID) Timer timer,
+        @TimerId(END_OF_WINDOW_ID) Timer windowTimer,
+        @TimerId(END_OF_BUFFERING_ID) Timer bufferingTimer,
         @StateId(BATCH_ID) BagState<InputT> batch,
         @StateId(NUM_ELEMENTS_IN_BATCH_ID) CombiningState<Long, long[], Long> 
numElementsInBatch,
         @StateId(KEY_ID) ValueState<K> key,
         @Element KV<K, InputT> element,
         BoundedWindow window,
         OutputReceiver<KV<K, Iterable<InputT>>> receiver) {
-      Instant windowExpires = window.maxTimestamp().plus(allowedLateness);
-
-      LOG.debug(
-          "*** SET TIMER *** to point in time {} for window {}",
-          windowExpires.toString(),
-          window.toString());
-      timer.set(windowExpires);
+      Instant windowEnds = window.maxTimestamp().plus(allowedLateness);
+      LOG.debug("*** SET TIMER *** to point in time {} for window {}", 
windowEnds, window);
+      windowTimer.set(windowEnds);
       key.write(element.getKey());
+      LOG.debug("*** BATCH *** Add element for window {} ", window);
       batch.add(element.getValue());
-      LOG.debug("*** BATCH *** Add element for window {} ", window.toString());
-      // blind add is supported with combiningState
+      // Blind add is supported with combiningState
       numElementsInBatch.add(1L);
+
       Long num = numElementsInBatch.read();
+      if (num == 1 && maxBufferingDuration.isLongerThan(Duration.ZERO)) {

Review comment:
       If we `checkArgument(duration != Duration.ZERO)` when 
`withMaxBufferingDuration ()`, then we can drop the check here.

##########
File path: 
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupIntoBatches.java
##########
@@ -66,25 +68,49 @@
  *         }
  *     }}));
  * </pre>
+ *
+ * <p>Example 2 (batch unbounded input in a global window):
+ *
+ * <pre>{@code
+ * PCollection<KV<String, String>> unboundedInput = ...;
+ * long batchSize = 100L;
+ * Duration maxBufferingDuration = Duration.standardSeconds(10);
+ * PCollection<KV<String, Iterable<String>>> batched = unboundedInput
+ *     .apply(Window.<KV<String, String>>into(new GlobalWindows())
+ *         .triggering(Repeatedly.forever(AfterPane.elementCountAtLeast(1)))
+ *         .discardingFiredPanes())
+ *     .apply(GroupIntoBatches.<String, String>ofSize(batchSize)
+ *         .withMaxBufferingDuration(maxBufferingDuration));
+ * }</pre>
  */
 public class GroupIntoBatches<K, InputT>
     extends PTransform<PCollection<KV<K, InputT>>, PCollection<KV<K, 
Iterable<InputT>>>> {
 
   private final long batchSize;
+  private final Duration maxBufferingDuration;
 
-  private GroupIntoBatches(long batchSize) {
+  private GroupIntoBatches(long batchSize, Duration maxBufferingDuration) {
     this.batchSize = batchSize;
+    this.maxBufferingDuration = maxBufferingDuration;
   }
 
   public static <K, InputT> GroupIntoBatches<K, InputT> ofSize(long batchSize) 
{
-    return new GroupIntoBatches<>(batchSize);
+    return new GroupIntoBatches<>(batchSize, Duration.ZERO);
   }
 
   /** Returns the size of the batch. */
   public long getBatchSize() {
     return batchSize;
   }
 
+  /**
+   * Set a time limit (in processing time) on how long an incomplete batch of 
elements is allowed to
+   * be buffered. Once a batch is flushed to output, the timer is reset.
+   */
+  public GroupIntoBatches<K, InputT> withMaxBufferingDuration(Duration 
duration) {
+    return new GroupIntoBatches<>(batchSize, duration);

Review comment:
       If `Duration.ZERO` is not preferred, we can have a check here.

##########
File path: 
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupIntoBatches.java
##########
@@ -150,70 +185,96 @@ public long apply(long left, long right) {
               });
 
       this.keySpec = StateSpecs.value(inputKeyCoder);
-      // prefetch every 20% of batchSize elements. Do not prefetch if 
batchSize is too little
+      // Prefetch every 20% of batchSize elements. Do not prefetch if 
batchSize is too little
       this.prefetchFrequency = ((batchSize / 5) <= 1) ? Long.MAX_VALUE : 
(batchSize / 5);
     }
 
     @ProcessElement
     public void processElement(
-        @TimerId(END_OF_WINDOW_ID) Timer timer,
+        @TimerId(END_OF_WINDOW_ID) Timer windowTimer,
+        @TimerId(END_OF_BUFFERING_ID) Timer bufferingTimer,
         @StateId(BATCH_ID) BagState<InputT> batch,
         @StateId(NUM_ELEMENTS_IN_BATCH_ID) CombiningState<Long, long[], Long> 
numElementsInBatch,
         @StateId(KEY_ID) ValueState<K> key,
         @Element KV<K, InputT> element,
         BoundedWindow window,
         OutputReceiver<KV<K, Iterable<InputT>>> receiver) {
-      Instant windowExpires = window.maxTimestamp().plus(allowedLateness);
-
-      LOG.debug(
-          "*** SET TIMER *** to point in time {} for window {}",
-          windowExpires.toString(),
-          window.toString());
-      timer.set(windowExpires);
+      Instant windowEnds = window.maxTimestamp().plus(allowedLateness);
+      LOG.debug("*** SET TIMER *** to point in time {} for window {}", 
windowEnds, window);
+      windowTimer.set(windowEnds);
       key.write(element.getKey());
+      LOG.debug("*** BATCH *** Add element for window {} ", window);
       batch.add(element.getValue());
-      LOG.debug("*** BATCH *** Add element for window {} ", window.toString());
-      // blind add is supported with combiningState
+      // Blind add is supported with combiningState
       numElementsInBatch.add(1L);
+
       Long num = numElementsInBatch.read();
+      if (num == 1 && maxBufferingDuration.isLongerThan(Duration.ZERO)) {
+        // This is the first element in batch. Start counting buffering time 
if a limit was set.
+        bufferingTimer.offset(maxBufferingDuration).setRelative();
+      }
       if (num % prefetchFrequency == 0) {
-        // prefetch data and modify batch state (readLater() modifies this)
+        // Prefetch data and modify batch state (readLater() modifies this)
         batch.readLater();
       }
       if (num >= batchSize) {
         LOG.debug("*** END OF BATCH *** for window {}", window.toString());
-        flushBatch(receiver, key, batch, numElementsInBatch);
+        flushBatch(receiver, key, batch, numElementsInBatch, bufferingTimer);
       }
     }
 
+    @OnTimer(END_OF_BUFFERING_ID)
+    public void onBufferingTimer(
+        OutputReceiver<KV<K, Iterable<InputT>>> receiver,
+        @Timestamp Instant timestamp,
+        @StateId(KEY_ID) ValueState<K> key,
+        @StateId(BATCH_ID) BagState<InputT> batch,
+        @StateId(NUM_ELEMENTS_IN_BATCH_ID) CombiningState<Long, long[], Long> 
numElementsInBatch,
+        @TimerId(END_OF_BUFFERING_ID) Timer bufferingTimer) {
+      LOG.debug(
+          "*** END OF BUFFERING *** for timer timestamp {} with buffering 
duration {}",
+          timestamp,
+          maxBufferingDuration);
+      flushBatch(receiver, key, batch, numElementsInBatch, bufferingTimer);
+    }
+
     @OnTimer(END_OF_WINDOW_ID)
-    public void onTimerCallback(
+    public void onWindowTimer(
         OutputReceiver<KV<K, Iterable<InputT>>> receiver,
         @Timestamp Instant timestamp,
         @StateId(KEY_ID) ValueState<K> key,
         @StateId(BATCH_ID) BagState<InputT> batch,
         @StateId(NUM_ELEMENTS_IN_BATCH_ID) CombiningState<Long, long[], Long> 
numElementsInBatch,
+        @TimerId(END_OF_BUFFERING_ID) Timer bufferingTimer,
         BoundedWindow window) {
       LOG.debug(
           "*** END OF WINDOW *** for timer timestamp {} in windows {}",
           timestamp,
           window.toString());
-      flushBatch(receiver, key, batch, numElementsInBatch);
+      flushBatch(receiver, key, batch, numElementsInBatch, null);
     }
 
     private void flushBatch(
         OutputReceiver<KV<K, Iterable<InputT>>> receiver,
         ValueState<K> key,
         BagState<InputT> batch,
-        CombiningState<Long, long[], Long> numElementsInBatch) {
+        CombiningState<Long, long[], Long> numElementsInBatch,
+        @Nullable Timer bufferingTimer) {
       Iterable<InputT> values = batch.read();
-      // when the timer fires, batch state might be empty
+      // When the timer fires, batch state might be empty
       if (!Iterables.isEmpty(values)) {
         receiver.output(KV.of(key.read(), values));
       }
       batch.clear();
       LOG.debug("*** BATCH *** clear");
       numElementsInBatch.clear();
+      // We might reach here due to batch size being reached or window 
expiration. Reset the
+      // buffering timer (if not null) since the state is empty now. It'll be 
extended again if a
+      // new element arrives prior to the expiration time set here.
+      // TODO(BEAM-10887): Use clear() when it's available.
+      if (bufferingTimer != null && 
maxBufferingDuration.isLongerThan(Duration.ZERO)) {

Review comment:
       Same above.

##########
File path: 
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupIntoBatches.java
##########
@@ -150,70 +185,96 @@ public long apply(long left, long right) {
               });
 
       this.keySpec = StateSpecs.value(inputKeyCoder);
-      // prefetch every 20% of batchSize elements. Do not prefetch if 
batchSize is too little
+      // Prefetch every 20% of batchSize elements. Do not prefetch if 
batchSize is too little
       this.prefetchFrequency = ((batchSize / 5) <= 1) ? Long.MAX_VALUE : 
(batchSize / 5);
     }
 
     @ProcessElement
     public void processElement(
-        @TimerId(END_OF_WINDOW_ID) Timer timer,
+        @TimerId(END_OF_WINDOW_ID) Timer windowTimer,
+        @TimerId(END_OF_BUFFERING_ID) Timer bufferingTimer,
         @StateId(BATCH_ID) BagState<InputT> batch,
         @StateId(NUM_ELEMENTS_IN_BATCH_ID) CombiningState<Long, long[], Long> 
numElementsInBatch,
         @StateId(KEY_ID) ValueState<K> key,
         @Element KV<K, InputT> element,
         BoundedWindow window,
         OutputReceiver<KV<K, Iterable<InputT>>> receiver) {
-      Instant windowExpires = window.maxTimestamp().plus(allowedLateness);
-
-      LOG.debug(
-          "*** SET TIMER *** to point in time {} for window {}",
-          windowExpires.toString(),
-          window.toString());
-      timer.set(windowExpires);
+      Instant windowEnds = window.maxTimestamp().plus(allowedLateness);
+      LOG.debug("*** SET TIMER *** to point in time {} for window {}", 
windowEnds, window);
+      windowTimer.set(windowEnds);
       key.write(element.getKey());
+      LOG.debug("*** BATCH *** Add element for window {} ", window);
       batch.add(element.getValue());
-      LOG.debug("*** BATCH *** Add element for window {} ", window.toString());
-      // blind add is supported with combiningState
+      // Blind add is supported with combiningState
       numElementsInBatch.add(1L);
+
       Long num = numElementsInBatch.read();
+      if (num == 1 && maxBufferingDuration.isLongerThan(Duration.ZERO)) {
+        // This is the first element in batch. Start counting buffering time 
if a limit was set.
+        bufferingTimer.offset(maxBufferingDuration).setRelative();
+      }
       if (num % prefetchFrequency == 0) {
-        // prefetch data and modify batch state (readLater() modifies this)
+        // Prefetch data and modify batch state (readLater() modifies this)
         batch.readLater();
       }
       if (num >= batchSize) {
         LOG.debug("*** END OF BATCH *** for window {}", window.toString());
-        flushBatch(receiver, key, batch, numElementsInBatch);
+        flushBatch(receiver, key, batch, numElementsInBatch, bufferingTimer);
       }
     }
 
+    @OnTimer(END_OF_BUFFERING_ID)
+    public void onBufferingTimer(
+        OutputReceiver<KV<K, Iterable<InputT>>> receiver,
+        @Timestamp Instant timestamp,
+        @StateId(KEY_ID) ValueState<K> key,
+        @StateId(BATCH_ID) BagState<InputT> batch,
+        @StateId(NUM_ELEMENTS_IN_BATCH_ID) CombiningState<Long, long[], Long> 
numElementsInBatch,
+        @TimerId(END_OF_BUFFERING_ID) Timer bufferingTimer) {
+      LOG.debug(
+          "*** END OF BUFFERING *** for timer timestamp {} with buffering 
duration {}",
+          timestamp,
+          maxBufferingDuration);
+      flushBatch(receiver, key, batch, numElementsInBatch, bufferingTimer);
+    }
+
     @OnTimer(END_OF_WINDOW_ID)
-    public void onTimerCallback(
+    public void onWindowTimer(
         OutputReceiver<KV<K, Iterable<InputT>>> receiver,
         @Timestamp Instant timestamp,
         @StateId(KEY_ID) ValueState<K> key,
         @StateId(BATCH_ID) BagState<InputT> batch,
         @StateId(NUM_ELEMENTS_IN_BATCH_ID) CombiningState<Long, long[], Long> 
numElementsInBatch,
+        @TimerId(END_OF_BUFFERING_ID) Timer bufferingTimer,
         BoundedWindow window) {
       LOG.debug(
           "*** END OF WINDOW *** for timer timestamp {} in windows {}",
           timestamp,
           window.toString());
-      flushBatch(receiver, key, batch, numElementsInBatch);
+      flushBatch(receiver, key, batch, numElementsInBatch, null);

Review comment:
       Instead, we want to pass actual bufferingTimer here since we want to 
clear it.




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


Reply via email to