From cb5d4274dcdae577ddffadbe48696ee001cd7c6f Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Fri, 23 Oct 2026 10:41:13 -0700 Subject: [PATCH 1/9] Fix GroupIntoBatches watermark hold update, timerTs state, early flush timestamp contamination, and parameter checks --- .../beam/sdk/transforms/GroupIntoBatches.java | 112 +++++++++++++----- .../sdk/transforms/GroupIntoBatchesTest.java | 75 ++++++++++++ 2 files changed, 158 insertions(+), 29 deletions(-) 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..878cc731ba91 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 @@ -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>> batched = input * .apply(GroupIntoBatches.ofSize(batchSize)) * .setCoder(KvCoder.of(StringUtf8Coder.of(), IterableCoder.of(StringUtf8Coder.of()))) - * .apply(ParDo.of(new DoFn>, KV>() }{ - * {@code @ProcessElement + * .apply(ParDo.of(new DoFn>, KV>() { + * @ProcessElement * public void processElement(@Element KV> element, * OutputReceiver> r) { * r.output(KV.of(element.getKey(), callWebService(element.getValue()))); * } - * }})); + * })); * * *

Example 2 (batch unbounded input in a global window): @@ -171,7 +170,9 @@ private GroupIntoBatches(BatchingParams params) { /** Aim to create batches each with the specified element count. */ public static GroupIntoBatches ofSize(long batchSize) { - Preconditions.checkState(batchSize < Long.MAX_VALUE); + checkArgument( + batchSize > 0 && batchSize < Long.MAX_VALUE, + "batchSize must be greater than 0 and less than Long.MAX_VALUE"); return new GroupIntoBatches(BatchingParams.createDefault()).withSize(batchSize); } @@ -210,9 +211,25 @@ 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 withSize(long batchSize) { - Preconditions.checkState(batchSize < Long.MAX_VALUE); + checkArgument( + batchSize > 0 && batchSize < Long.MAX_VALUE, + "batchSize must be greater than 0 and less than Long.MAX_VALUE"); return new GroupIntoBatches<>( BatchingParams.create( batchSize, @@ -223,7 +240,9 @@ public GroupIntoBatches withSize(long batchSize) { /** See {@link #ofByteSize(long)}. */ public GroupIntoBatches withByteSize(long batchSizeBytes) { - Preconditions.checkState(batchSizeBytes < Long.MAX_VALUE); + checkArgument( + batchSizeBytes > 0 && batchSizeBytes < Long.MAX_VALUE, + "batchSizeBytes must be greater than 0 and less than Long.MAX_VALUE"); return new GroupIntoBatches<>( BatchingParams.create( params.getBatchSize(), @@ -235,7 +254,9 @@ public GroupIntoBatches withByteSize(long batchSizeBytes) { /** See {@link #ofByteSize(long, SerializableFunction)}. */ public GroupIntoBatches withByteSize( long batchSizeBytes, SerializableFunction getElementByteSize) { - Preconditions.checkState(batchSizeBytes < Long.MAX_VALUE); + checkArgument( + batchSizeBytes > 0 && batchSizeBytes < Long.MAX_VALUE, + "batchSizeBytes must be greater than 0 and less than Long.MAX_VALUE"); return new GroupIntoBatches<>( BatchingParams.create( params.getBatchSize(), @@ -308,6 +329,20 @@ public KV, InputT> apply(KV 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 +357,6 @@ public long getElementByteSize() { return this.elementByteSize; } } - ; @Override public PCollection>> expand(PCollection> input) { @@ -457,7 +491,10 @@ public long apply(long left, long right) { this.minBufferedTsSpec = StateSpecs.combining(minCombineFn); // Prefetch every 20% of batchSize elements. Do not prefetch if batchSize is too little - this.prefetchFrequency = ((batchSize / 5) <= 1) ? Long.MAX_VALUE : (batchSize / 5); + this.prefetchFrequency = + (batchSize == Long.MAX_VALUE) + ? 100L + : ((batchSize / 5) <= 1) ? Long.MAX_VALUE : (batchSize / 5); } @Override @@ -492,13 +529,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 +562,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); } } @@ -619,8 +665,16 @@ public void onWindowExpiration( } @OnTimer(TIMER_HOLD_ID) - public void onHoldTimer() { - // Do nothing. The associated watermark hold will be automatically removed. + public void onHoldTimer( + OutputReceiver>> receiver, + @Key K key, + @StateId(BATCH_ID) BagState batch, + @StateId(NUM_ELEMENTS_IN_BATCH_ID) CombiningState storedBatchSize, + @StateId(NUM_BYTES_IN_BATCH_ID) CombiningState storedBatchSizeBytes, + @StateId(TIMER_TIMESTAMP) ValueState timerTs, + @StateId(MIN_BUFFERED_TS) CombiningState minBufferedTs) { + flushBatch( + receiver, key, batch, storedBatchSize, storedBatchSizeBytes, timerTs, minBufferedTs); } // We no longer set this timer, since OnWindowExpiration takes care of his. However we leave the 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..b09733e0f04e 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 @@ -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,76 @@ private static boolean checkBatchByteSizes( } return true; } + + @Test(expected = IllegalArgumentException.class) + public void testZeroBatchSizeThrows() { + GroupIntoBatches.ofSize(0); + } + + @Test(expected = IllegalArgumentException.class) + public void testNegativeBatchSizeThrows() { + GroupIntoBatches.ofSize(-1); + } + + @Test(expected = IllegalArgumentException.class) + public void testZeroBatchSizeBytesThrows() { + GroupIntoBatches.ofByteSize(0); + } + + @Test(expected = IllegalArgumentException.class) + public void testNegativeBatchSizeBytesThrows() { + GroupIntoBatches.ofByteSize(-1); + } + + @Test + public void testDisplayData() { + GroupIntoBatches gib = + GroupIntoBatches.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))); + } + + @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(5)))) + .advanceWatermarkToInfinity(); + + PCollection outputTimestamps = + pipeline + .apply(stream) + .apply(GroupIntoBatches.ofSize(2)) + .apply( + ParDo.of( + new DoFn>, Long>() { + @ProcessElement + public void processElement( + @Timestamp Instant ts, OutputReceiver receiver) { + receiver.output(ts.getMillis()); + } + })); + + PAssert.that(outputTimestamps) + .containsInAnyOrder(baseTs.plus(Duration.standardSeconds(5)).getMillis()); + pipeline.run(); + } } From b2bcbd7b812574eda1e17448ec08ed2a0ca6a793 Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Fri, 23 Oct 2026 10:46:22 -0700 Subject: [PATCH 2/9] Allow non-positive batchSize in GroupIntoBatches to maintain backwards compatibility --- .../beam/sdk/transforms/GroupIntoBatches.java | 14 ++++---------- .../sdk/transforms/GroupIntoBatchesTest.java | 18 ++++-------------- 2 files changed, 8 insertions(+), 24 deletions(-) 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 878cc731ba91..372fe2129ded 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 @@ -170,9 +170,7 @@ private GroupIntoBatches(BatchingParams params) { /** Aim to create batches each with the specified element count. */ public static GroupIntoBatches ofSize(long batchSize) { - checkArgument( - batchSize > 0 && batchSize < Long.MAX_VALUE, - "batchSize must be greater than 0 and less than Long.MAX_VALUE"); + checkArgument(batchSize < Long.MAX_VALUE, "batchSize must be less than Long.MAX_VALUE"); return new GroupIntoBatches(BatchingParams.createDefault()).withSize(batchSize); } @@ -227,9 +225,7 @@ public void populateDisplayData(DisplayData.Builder builder) { /** See {@link #ofSize(long)}. */ public GroupIntoBatches withSize(long batchSize) { - checkArgument( - batchSize > 0 && batchSize < Long.MAX_VALUE, - "batchSize must be greater than 0 and less than Long.MAX_VALUE"); + checkArgument(batchSize < Long.MAX_VALUE, "batchSize must be less than Long.MAX_VALUE"); return new GroupIntoBatches<>( BatchingParams.create( batchSize, @@ -241,8 +237,7 @@ public GroupIntoBatches withSize(long batchSize) { /** See {@link #ofByteSize(long)}. */ public GroupIntoBatches withByteSize(long batchSizeBytes) { checkArgument( - batchSizeBytes > 0 && batchSizeBytes < Long.MAX_VALUE, - "batchSizeBytes must be greater than 0 and less than Long.MAX_VALUE"); + batchSizeBytes < Long.MAX_VALUE, "batchSizeBytes must be less than Long.MAX_VALUE"); return new GroupIntoBatches<>( BatchingParams.create( params.getBatchSize(), @@ -255,8 +250,7 @@ public GroupIntoBatches withByteSize(long batchSizeBytes) { public GroupIntoBatches withByteSize( long batchSizeBytes, SerializableFunction getElementByteSize) { checkArgument( - batchSizeBytes > 0 && batchSizeBytes < Long.MAX_VALUE, - "batchSizeBytes must be greater than 0 and less than Long.MAX_VALUE"); + batchSizeBytes < Long.MAX_VALUE, "batchSizeBytes must be less than Long.MAX_VALUE"); return new GroupIntoBatches<>( BatchingParams.create( params.getBatchSize(), 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 b09733e0f04e..50a996d72db9 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 @@ -810,23 +810,13 @@ private static boolean checkBatchByteSizes( } @Test(expected = IllegalArgumentException.class) - public void testZeroBatchSizeThrows() { - GroupIntoBatches.ofSize(0); + public void testMaxBatchSizeThrows() { + GroupIntoBatches.ofSize(Long.MAX_VALUE); } @Test(expected = IllegalArgumentException.class) - public void testNegativeBatchSizeThrows() { - GroupIntoBatches.ofSize(-1); - } - - @Test(expected = IllegalArgumentException.class) - public void testZeroBatchSizeBytesThrows() { - GroupIntoBatches.ofByteSize(0); - } - - @Test(expected = IllegalArgumentException.class) - public void testNegativeBatchSizeBytesThrows() { - GroupIntoBatches.ofByteSize(-1); + public void testMaxBatchSizeBytesThrows() { + GroupIntoBatches.ofByteSize(Long.MAX_VALUE); } @Test From 3f026b82b2142e13c3335b01485715d024aea4db Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Fri, 23 Oct 2026 10:48:13 -0700 Subject: [PATCH 3/9] Revert prefetchFrequency change for batchSize == Long.MAX_VALUE --- .../org/apache/beam/sdk/transforms/GroupIntoBatches.java | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) 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 372fe2129ded..ce7f513db22c 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 @@ -485,10 +485,7 @@ public long apply(long left, long right) { this.minBufferedTsSpec = StateSpecs.combining(minCombineFn); // Prefetch every 20% of batchSize elements. Do not prefetch if batchSize is too little - this.prefetchFrequency = - (batchSize == Long.MAX_VALUE) - ? 100L - : ((batchSize / 5) <= 1) ? Long.MAX_VALUE : (batchSize / 5); + this.prefetchFrequency = ((batchSize / 5) <= 1) ? Long.MAX_VALUE : (batchSize / 5); } @Override From 6e008e54926e6dfb54853f865c73f9eab3b7eb09 Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Fri, 23 Oct 2026 11:12:39 -0700 Subject: [PATCH 4/9] Revert onHoldTimer callback to original no-op implementation --- .../apache/beam/sdk/transforms/GroupIntoBatches.java | 12 ++---------- 1 file changed, 2 insertions(+), 10 deletions(-) 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 ce7f513db22c..970b92592d6e 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 @@ -656,16 +656,8 @@ public void onWindowExpiration( } @OnTimer(TIMER_HOLD_ID) - public void onHoldTimer( - OutputReceiver>> receiver, - @Key K key, - @StateId(BATCH_ID) BagState batch, - @StateId(NUM_ELEMENTS_IN_BATCH_ID) CombiningState storedBatchSize, - @StateId(NUM_BYTES_IN_BATCH_ID) CombiningState storedBatchSizeBytes, - @StateId(TIMER_TIMESTAMP) ValueState timerTs, - @StateId(MIN_BUFFERED_TS) CombiningState minBufferedTs) { - flushBatch( - receiver, key, batch, storedBatchSize, storedBatchSizeBytes, timerTs, minBufferedTs); + public void onHoldTimer() { + // Do nothing. The associated watermark hold will be automatically removed. } // We no longer set this timer, since OnWindowExpiration takes care of his. However we leave the From 0fd63989ee36605da5c18be61135c97675e02e65 Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Fri, 23 Oct 2026 11:29:10 -0700 Subject: [PATCH 5/9] Update testOutOfOrderTimestampsInBatch to ofSize(3) to test watermark hold on window expiration --- .../apache/beam/sdk/transforms/GroupIntoBatchesTest.java | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) 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 50a996d72db9..054ca24d7a53 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 @@ -831,6 +831,11 @@ public void testDisplayData() { 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, @@ -855,7 +860,7 @@ public void testOutOfOrderTimestampsInBatch() { PCollection outputTimestamps = pipeline .apply(stream) - .apply(GroupIntoBatches.ofSize(2)) + .apply(GroupIntoBatches.ofSize(3)) .apply( ParDo.of( new DoFn>, Long>() { From bab7761af9899b321244a67eb23e0db963e95cde Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Mon, 26 Oct 2026 10:31:25 -0700 Subject: [PATCH 6/9] Revert changes unrelated to watermark hold for out of order elements --- .../beam/sdk/transforms/GroupIntoBatches.java | 55 ++++--------------- .../sdk/transforms/GroupIntoBatchesTest.java | 25 --------- 2 files changed, 11 insertions(+), 69 deletions(-) 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 970b92592d6e..33618e22e02a 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 @@ -43,6 +43,7 @@ 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; @@ -77,13 +78,13 @@ * PCollection>> batched = input * .apply(GroupIntoBatches.ofSize(batchSize)) * .setCoder(KvCoder.of(StringUtf8Coder.of(), IterableCoder.of(StringUtf8Coder.of()))) - * .apply(ParDo.of(new DoFn>, KV>() { - * @ProcessElement + * .apply(ParDo.of(new DoFn>, KV>() }{ + * {@code @ProcessElement * public void processElement(@Element KV> element, * OutputReceiver> r) { * r.output(KV.of(element.getKey(), callWebService(element.getValue()))); * } - * })); + * }})); * * *

Example 2 (batch unbounded input in a global window): @@ -170,7 +171,7 @@ private GroupIntoBatches(BatchingParams params) { /** Aim to create batches each with the specified element count. */ public static GroupIntoBatches ofSize(long batchSize) { - checkArgument(batchSize < Long.MAX_VALUE, "batchSize must be less than Long.MAX_VALUE"); + Preconditions.checkState(batchSize < Long.MAX_VALUE); return new GroupIntoBatches(BatchingParams.createDefault()).withSize(batchSize); } @@ -209,23 +210,9 @@ 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 withSize(long batchSize) { - checkArgument(batchSize < Long.MAX_VALUE, "batchSize must be less than Long.MAX_VALUE"); + Preconditions.checkState(batchSize < Long.MAX_VALUE); return new GroupIntoBatches<>( BatchingParams.create( batchSize, @@ -236,8 +223,7 @@ public GroupIntoBatches withSize(long batchSize) { /** See {@link #ofByteSize(long)}. */ public GroupIntoBatches withByteSize(long batchSizeBytes) { - checkArgument( - batchSizeBytes < Long.MAX_VALUE, "batchSizeBytes must be less than Long.MAX_VALUE"); + Preconditions.checkState(batchSizeBytes < Long.MAX_VALUE); return new GroupIntoBatches<>( BatchingParams.create( params.getBatchSize(), @@ -249,8 +235,7 @@ public GroupIntoBatches withByteSize(long batchSizeBytes) { /** See {@link #ofByteSize(long, SerializableFunction)}. */ public GroupIntoBatches withByteSize( long batchSizeBytes, SerializableFunction getElementByteSize) { - checkArgument( - batchSizeBytes < Long.MAX_VALUE, "batchSizeBytes must be less than Long.MAX_VALUE"); + Preconditions.checkState(batchSizeBytes < Long.MAX_VALUE); return new GroupIntoBatches<>( BatchingParams.create( params.getBatchSize(), @@ -323,20 +308,6 @@ public KV, InputT> apply(KV 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 { @@ -351,6 +322,7 @@ public long getElementByteSize() { return this.elementByteSize; } } + ; @Override public PCollection>> expand(PCollection> input) { @@ -523,8 +495,7 @@ public void processElement( LOG.debug("*** BATCH *** Add element for window {} ", window); if (shouldCareAboutWeight) { final long elementWeight = weigher.apply(element.getValue()); - if (storedBatchSizeBytes.read() > 0 - && elementWeight + storedBatchSizeBytes.read() > batchSizeBytes) { + if (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 @@ -588,11 +559,7 @@ public void processElement( // 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 windowMax = window.maxTimestamp(); - Instant windowEnd = - windowMax.isBefore(BoundedWindow.TIMESTAMP_MAX_VALUE.minus(allowedLateness)) - ? windowMax.plus(allowedLateness) - : BoundedWindow.TIMESTAMP_MAX_VALUE; + Instant windowEnd = window.maxTimestamp().plus(allowedLateness); 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 054ca24d7a53..97ac04f871d1 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 @@ -17,8 +17,6 @@ */ 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; @@ -46,7 +44,6 @@ 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; @@ -809,28 +806,6 @@ private static 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 gib = - GroupIntoBatches.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, From f6660d95eea5e14063adcaaa378525a86ef81e93 Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Mon, 26 Oct 2026 10:51:34 -0700 Subject: [PATCH 7/9] Revert timerTs persistence changes --- .../beam/sdk/transforms/GroupIntoBatches.java | 24 +++++++------------ 1 file changed, 9 insertions(+), 15 deletions(-) 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 33618e22e02a..d070125fca4d 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 @@ -525,7 +525,7 @@ public void processElement( storedBatchSizeBytes.add(elementWeight); } - long oldMinBufferedTs = + long oldOutputTs = MoreObjects.firstNonNull( minBufferedTs.read(), BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis()); @@ -535,32 +535,26 @@ public void processElement( 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. - boolean needsNewTimer = num == 1 || currentMinBufferedTs < oldMinBufferedTs; + boolean needsNewTimer = num == 1 || minBufferedTs.read() < oldOutputTs; if (needsNewTimer) { if (shouldCareAboutMaxBufferingDuration) { - Long existingTargetTs = timerTs.read(); - long targetTs; - if (existingTargetTs == null) { - targetTs = - bufferingTimer.getCurrentRelativeTime().getMillis() - + maxBufferingDuration.getMillis(); - timerTs.write(targetTs); - } else { - targetTs = existingTargetTs; - } + long targetTs = + MoreObjects.firstNonNull( + timerTs.read(), + bufferingTimer.getCurrentRelativeTime().getMillis() + + maxBufferingDuration.getMillis()); bufferingTimer - .withOutputTimestamp(Instant.ofEpochMilli(currentMinBufferedTs)) + .withOutputTimestamp(Instant.ofEpochMilli(minBufferedTs.read())) .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(currentMinBufferedTs)).set(windowEnd); + holdTimer.withOutputTimestamp(Instant.ofEpochMilli(minBufferedTs.read())).set(windowEnd); } } From dd1527b96f6f6867efcb51bfe4f9d68a755a87ba Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Mon, 26 Oct 2026 10:55:23 -0700 Subject: [PATCH 8/9] Reuse currentMinBufferedTs to avoid multiple minBufferedTs.read() calls --- .../org/apache/beam/sdk/transforms/GroupIntoBatches.java | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) 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 d070125fca4d..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 @@ -535,10 +535,11 @@ public void processElement( 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. - boolean needsNewTimer = num == 1 || minBufferedTs.read() < oldOutputTs; + boolean needsNewTimer = num == 1 || currentMinBufferedTs < oldOutputTs; if (needsNewTimer) { if (shouldCareAboutMaxBufferingDuration) { long targetTs = @@ -547,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); } } From ca981ea12c07a890fb9ddbd8cef19e2fcc43be61 Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Mon, 26 Oct 2026 16:10:29 -0700 Subject: [PATCH 9/9] Strengthen testOutOfOrderTimestampsInBatch to verify watermark hold prevents downstream late data drops --- .../sdk/transforms/GroupIntoBatchesTest.java | 32 +++++++++++++------ 1 file changed, 22 insertions(+), 10 deletions(-) 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 97ac04f871d1..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 @@ -808,8 +808,10 @@ private static boolean checkBatchByteSizes( /** * 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. + * 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({ @@ -829,25 +831,35 @@ public void testOutOfOrderTimestampsInBatch() { TimestampedValue.of( KV.of("key", "val1"), baseTs.plus(Duration.standardSeconds(10)))) .addElements( - TimestampedValue.of(KV.of("key", "val2"), baseTs.plus(Duration.standardSeconds(5)))) + TimestampedValue.of(KV.of("key", "val2"), baseTs.plus(Duration.standardSeconds(2)))) + .advanceWatermarkTo(baseTs.plus(Duration.standardSeconds(7))) .advanceWatermarkToInfinity(); - PCollection outputTimestamps = + PCollection>> output = pipeline .apply(stream) .apply(GroupIntoBatches.ofSize(3)) + .apply(Window.into(FixedWindows.of(Duration.standardSeconds(5)))) .apply( ParDo.of( - new DoFn>, Long>() { + new DoFn>, KV>() { @ProcessElement public void processElement( - @Timestamp Instant ts, OutputReceiver receiver) { - receiver.output(ts.getMillis()); + @Element KV> element, + OutputReceiver> receiver) { + receiver.output(KV.of(element.getKey(), "val")); } - })); + })) + .apply(GroupByKey.create()); - PAssert.that(outputTimestamps) - .containsInAnyOrder(baseTs.plus(Duration.standardSeconds(5)).getMillis()); + PAssert.that(output) + .satisfies( + iterable -> { + assertTrue( + "Expected output from GroupByKey but was dropped as late data", + iterable.iterator().hasNext()); + return null; + }); pipeline.run(); } }