[GitHub] [beam] boyuanzz commented on a change in pull request #12726: [BEAM-10475] Add max buffering duration to GroupIntoBatches (Java)

[GitBox <gi...@apache.org>]

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



##########
File path: sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupIntoBatches.java
##########
@@ -150,70 +184,95 @@ 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, bufferingTimer);
     }
 
     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,
+        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, window expiration or buffering time
+      // limit being reached. Reset the buffering timer anyway since the state is empty now. It'll
+      // be reset again if a new element arrives before the expiration time set here.
+      if (maxBufferingDuration.isLongerThan(Duration.ZERO)) {
+        bufferingTimer.offset(maxBufferingDuration).setRelative();

Review comment:
       It seems like we want to clear the timer when the flush is triggered by batch size/end of window. I think we can just clear the timer when `flushBatch` is called by window timer expiration and `num >= batchSize`.




----------------------------------------------------------------
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:
users@infra.apache.org