diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/CountingSource.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/CountingSource.java index 9d30efb2f113..4ef928480d75 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/CountingSource.java +++ b/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 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 unbounded() { @Deprecated public static UnboundedSource unboundedWithTimestampFn( SerializableFunction 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 { + public static class UnboundedCountingSource extends UnboundedSource { /** 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 private UnboundedCountingSource( long start, long stride, + long end, long elementsPerPeriod, Duration period, SerializableFunction 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 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> 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,6 +394,7 @@ 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); @@ -383,7 +402,7 @@ public boolean equals(@Nullable Object other) { @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 diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/CountingSourceTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/CountingSourceTest.java index 70a09083619d..337462340df1 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/CountingSourceTest.java +++ b/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 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 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)))); + } }