Stop emitting upon truncate in Java PeriodicSequence (#25716)

apache · Mar 4, 2023 · 4da6025 · 4da6025
1 parent 39cab94
commit 4da6025
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/PeriodicSequence.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/PeriodicSequence.java
@@ -194,6 +194,12 @@ public WatermarkEstimator<Instant> newWatermarkEstimator(
       return new WatermarkEstimators.Manual(state);
     }
 
+    @TruncateRestriction
+    public RestrictionTracker.TruncateResult<OffsetRange> truncate() {
+      // stop emitting immediately upon drain
+      return null;
+    }
+
     @ProcessElement
     public ProcessContinuation processElement(
         @Element SequenceDefinition srcElement,
@@ -207,7 +213,7 @@ public ProcessContinuation processElement(
 
       boolean claimSuccess = true;
 
-      estimator.setWatermark(Instant.ofEpochMilli(restriction.getFrom()));
+      estimator.setWatermark(Instant.ofEpochMilli(nextOutput));
 
       while (claimSuccess && Instant.ofEpochMilli(nextOutput).isBeforeNow()) {
         claimSuccess = restrictionTracker.tryClaim(nextOutput);