diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupIntoBatches.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupIntoBatches.java index 74930bebad37..602414bcd7fb 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupIntoBatches.java +++ b/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); } } diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/GroupIntoBatchesTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/GroupIntoBatchesTest.java index 92cb006c6f96..632aaacdbd0f 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/GroupIntoBatchesTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/GroupIntoBatchesTest.java @@ -805,4 +805,61 @@ private static 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> 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>> output = + pipeline + .apply(stream) + .apply(GroupIntoBatches.ofSize(3)) + .apply(Window.into(FixedWindows.of(Duration.standardSeconds(5)))) + .apply( + ParDo.of( + new DoFn>, KV>() { + @ProcessElement + public void processElement( + @Element KV> element, + OutputReceiver> 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(); + } }