Add UnboundedCountingSource::to for bounded reads from an UnboundedCountingSource

sdks/java/core/src/main/java/org/apache/beam/sdk/io/CountingSource.java

@@ -39,6 +39,7 @@
 import org.apache.beam.sdk.metrics.SourceMetrics;
 import org.apache.beam.sdk.options.PipelineOptions;
 import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
 import org.checkerframework.checker.nullness.qual.Nullable;
@@ -95,7 +96,8 @@ static BoundedSource<Long> createSourceForSubrange(long startIndex, long endInde
   /** Create a new {@link UnboundedCountingSource}. */
   // package-private to return a typed UnboundedCountingSource rather than the UnboundedSource type.
   static UnboundedCountingSource createUnboundedFrom(long start) {
-    return new UnboundedCountingSource(start, 1, 1L, Duration.ZERO, new NowTimestampFn());
+    return new UnboundedCountingSource(
+        start, 1, Long.MAX_VALUE, 1L, Duration.ZERO, new NowTimestampFn());
   }
 
   /**
@@ -130,7 +132,7 @@ public static UnboundedSource<Long, CounterMark> unbounded() {
   @Deprecated
   public static UnboundedSource<Long, CounterMark> unboundedWithTimestampFn(
       SerializableFunction<Long, Instant> timestampFn) {
-    return new UnboundedCountingSource(0, 1, 1L, Duration.ZERO, timestampFn);
+    return new UnboundedCountingSource(0, 1, Long.MAX_VALUE, 1L, Duration.ZERO, timestampFn);
   }
 
   /////////////////////////////////////////////////////////////////////////////////////////////
@@ -267,11 +269,13 @@ public void close() throws IOException {}
   }
 
   /** An implementation of {@link CountingSource} that produces an unbounded {@link PCollection}. */
-  static class UnboundedCountingSource extends UnboundedSource<Long, CounterMark> {
+  public static class UnboundedCountingSource extends UnboundedSource<Long, CounterMark> {
     /** The first number (>= 0) generated by this {@link UnboundedCountingSource}. */
     private final long start;
     /** The interval between numbers generated by this {@link UnboundedCountingSource}. */
     private final long stride;
+    /** The exclusive limit for the sequence. */
+    private final long end;
     /** The number of elements to produce each period. */
     private final long elementsPerPeriod;
     /** The time between producing numbers from this {@link UnboundedCountingSource}. */
@@ -291,11 +295,13 @@ static class UnboundedCountingSource extends UnboundedSource<Long, CounterMark>
     private UnboundedCountingSource(
         long start,
         long stride,
+        long end,
         long elementsPerPeriod,
         Duration period,
         SerializableFunction<Long, Instant> timestampFn) {
       this.start = start;
       this.stride = stride;
+      this.end = end;
       checkArgument(
           elementsPerPeriod > 0L,
           "Must produce at least one element per period, got %s",
@@ -312,7 +318,8 @@ private UnboundedCountingSource(
      * will be produced with an interval between them equal to the period.
      */
     public UnboundedCountingSource withRate(long elementsPerPeriod, Duration period) {
-      return new UnboundedCountingSource(start, stride, elementsPerPeriod, period, timestampFn);
+      return new UnboundedCountingSource(
+          start, stride, end, elementsPerPeriod, period, timestampFn);
     }
 
     /**
@@ -324,7 +331,18 @@ public UnboundedCountingSource withRate(long elementsPerPeriod, Duration period)
     public UnboundedCountingSource withTimestampFn(
         SerializableFunction<Long, Instant> timestampFn) {
       checkNotNull(timestampFn);
-      return new UnboundedCountingSource(start, stride, elementsPerPeriod, period, timestampFn);
+      return new UnboundedCountingSource(
+          start, stride, end, elementsPerPeriod, period, timestampFn);
+    }
+
+    /**
+     * Returns an {@link UnboundedCountingSource} like this one but with the specified exclusive
+     * limit.
+     */
+    public UnboundedCountingSource to(long end) {
+      checkArgument(end >= start, "end (%s) must be >= start (%s)", end, start);
+      return new UnboundedCountingSource(
+          start, stride, end, elementsPerPeriod, period, timestampFn);
     }
 
     /**
@@ -348,7 +366,7 @@ public List<? extends UnboundedSource<Long, CountingSource.CounterMark>> split(
         // 0, 2, and 4.
         splits.add(
             new UnboundedCountingSource(
-                start + i * stride, newStride, elementsPerPeriod, period, timestampFn));
+                start + i * stride, newStride, end, elementsPerPeriod, period, timestampFn));
       }
       return splits.build();
     }
@@ -376,14 +394,15 @@ public boolean equals(@Nullable Object other) {
       UnboundedCountingSource that = (UnboundedCountingSource) other;
       return this.start == that.start
           && this.stride == that.stride
+          && this.end == that.end
           && this.elementsPerPeriod == that.elementsPerPeriod
           && Objects.equals(this.period, that.period)
           && Objects.equals(this.timestampFn, that.timestampFn);
     }
 
     @Override
     public int hashCode() {
-      return Objects.hash(start, stride, elementsPerPeriod, period, timestampFn);
+      return Objects.hash(start, stride, end, elementsPerPeriod, period, timestampFn);
     }
   }
 
@@ -431,6 +450,9 @@ public boolean advance() throws IOException {
         return false;
       }
       long nextValue = current + source.stride;
+      if (nextValue >= source.end) {
+        return false;
+      }
       if (expectedValue() < nextValue) {
         return false;
       }
@@ -453,7 +475,9 @@ private long expectedValue() {
 
     @Override
     public Instant getWatermark() {
-      return source.timestampFn.apply(current);
+      return (current >= source.end - source.stride)
+          ? BoundedWindow.TIMESTAMP_MAX_VALUE
+          : source.timestampFn.apply(current);
     }
 
     @Override

sdks/java/core/src/test/java/org/apache/beam/sdk/io/CountingSourceTest.java

@@ -21,6 +21,7 @@
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.lessThan;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
 import java.io.IOException;
@@ -315,4 +316,39 @@ public void testUnboundedSourceCheckpointMark() throws Exception {
     assertEquals(numToSkip + 1, (long) reader.getCurrent());
     assertEquals(numToSkip + 1, reader.getCurrentTimestamp().getMillis());
   }
+
+  @Test
+  @Category(NeedsRunner.class)
+  public void testUnboundedSourceWithFromAndTo() {
+    long start = 5;
+    long end = 10;
+    PCollection<Long> input = p.apply(Read.from(CountingSource.createUnboundedFrom(start).to(end)));
+
+    PAssert.that(input).containsInAnyOrder(5L, 6L, 7L, 8L, 9L);
+    p.run();
+  }
+
+  @Test
+  @Category(NeedsRunner.class)
+  public void testUnboundedSourceWithFromAndToAndRate() {
+    long start = 5;
+    long end = 15;
+    long elementsPerPeriod = 2;
+    Duration period = Duration.millis(10);
+    PCollection<Long> input =
+        p.apply(
+            Read.from(
+                CountingSource.createUnboundedFrom(start)
+                    .to(end)
+                    .withRate(elementsPerPeriod, period)));
+
+    PAssert.that(input).containsInAnyOrder(5L, 6L, 7L, 8L, 9L, 10L, 11L, 12L, 13L, 14L);
+
+    Instant startTime = Instant.now();
+    p.run();
+    Instant endTime = Instant.now();
+
+    long expectedMinimumMillis = ((end - start) * period.getMillis()) / elementsPerPeriod;
+    assertFalse(endTime.isBefore(startTime.plus(Duration.millis(expectedMinimumMillis))));
+  }
 }