This is an automated email from the ASF dual-hosted git repository.
scwhittle pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 0165344ab13 Add UnboundedCountingSource::to for bounded reads from an
UnboundedCountingSource (#39084)
0165344ab13 is described below
commit 0165344ab138084d2e2998bbc7129decf0350f0f
Author: Arun Pandian <[email protected]>
AuthorDate: Wed Jun 24 05:21:44 2026 -0700
Add UnboundedCountingSource::to for bounded reads from an
UnboundedCountingSource (#39084)
* Add UnboundedCountingSource::to for bounded reads from an
UnboundedCountingSource
UnboundedCountingSource currently does not support configuring an end
limit, this PR adds that capability.
GenerateSequence with a from + to + rate boils down to a
BoundedReadFromUnboundedSource. BoundedReadFromUnboundedSource splits
emit all outputs from a single bundle that does not work well on Streaming
jobs.
UnboundedCountingSource::to is currently not exposed outside
CountingSource. Eventually GenerateSequence can use
UnboundedCountingSource::to instead of BoundedReadFromUnboundedSource
---
.../org/apache/beam/sdk/io/CountingSource.java | 40 +++++++++++++++++-----
.../org/apache/beam/sdk/io/CountingSourceTest.java | 36 +++++++++++++++++++
2 files changed, 68 insertions(+), 8 deletions(-)
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 9d30efb2f11..4ef928480d7 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.Counter;
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 @@ public class CountingSource {
/** 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 class CountingSource {
@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 class CountingSource {
}
/** 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 @@ public class CountingSource {
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 @@ public class CountingSource {
* 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 class CountingSource {
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 class CountingSource {
// 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 class CountingSource {
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 class CountingSource {
@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 class CountingSource {
return false;
}
long nextValue = current + source.stride;
+ if (nextValue >= source.end) {
+ return false;
+ }
if (expectedValue() < nextValue) {
return false;
}
@@ -453,7 +475,9 @@ public class CountingSource {
@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 70a09083619..337462340df 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.MatcherAssert.assertThat;
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 class CountingSourceTest {
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))));
+ }
}
