withTimestampCursor(Duration allowedLatenes
return toBuilder().setTimestampCursorAllowedLateness(allowedLateness).build();
}
+ /**
+ * Limits how many previously unseen outputs from one poll are kept as checkpointed pending
+ * work.
+ *
+ * Passing {@code null} removes a previously configured limit and restores the default
+ * unlimited behavior. A non-null limit must be greater than zero.
+ *
+ *
If an incomplete poll contains more than {@code maxPendingResults} distinct unseen
+ * outputs, the oldest outputs are emitted first and the rest are omitted from that round. The
+ * poll's watermark is held at or before the omitted outputs so the function will be called
+ * again to retrieve them. A {@link PollResult#complete} is never limited because its contract
+ * promises that the function will not be called again.
+ *
+ *
This option must only be used with a {@link PollFn} that reports omitted outputs again in
+ * subsequent polls. Otherwise those outputs cannot be recovered. Without this option, all
+ * unseen outputs from each poll continue to be kept, preserving the existing behavior. This
+ * limits the persisted pending outputs, not the result list returned by {@link PollFn} or the
+ * temporary hashes used to preserve duplicate-key semantics while selecting the oldest outputs.
+ *
+ *
While a poll is truncated, its termination condition is deferred until a later poll is not
+ * truncated, so omitted outputs are not silently lost. This option is therefore intended for a
+ * snapshot-style {@link PollFn} that repeats a finite set of outputs until they are observed. A
+ * function that continuously returns more distinct unseen outputs than this limit can keep the
+ * input polling indefinitely even if its termination condition would otherwise stop it.
+ *
+ *
Combine this with {@link #withTimestampCursor()} to time-bound the completed-key state
+ * when output timestamps advance past the retention floor. It does not impose a count bound on
+ * completed keys that share a timestamp still at or above that floor.
+ */
+ public Growth withMaxPendingResults(
+ @Nullable Integer maxPendingResults) {
+ checkArgument(
+ maxPendingResults == null || maxPendingResults > 0,
+ "maxPendingResults must be greater than zero, but was %s",
+ maxPendingResults);
+ return toBuilder().setMaxPendingResults(maxPendingResults).build();
+ }
+
@Override
public PCollection> expand(PCollection input) {
checkNotNull(getPollInterval(), "pollInterval");
@@ -945,7 +993,9 @@ public ProcessContinuation process(
@Nullable Duration allowedLateness = spec.getTimestampCursorAllowedLateness();
@Nullable Instant cursor = pollingRestriction.getCursor();
- if (retentionFloorAtMaxTimestamp(cursor, allowedLateness)) {
+ if (retentionFloorAtMaxTimestamp(cursor, allowedLateness)
+ && (spec.getMaxPendingResults() == null
+ || BoundedWindow.TIMESTAMP_MAX_VALUE.equals(pollingRestriction.getPollWatermark()))) {
// Nothing can be claimed above the floor, so claim an empty round and stop.
LOG.info("{} - will not poll, retention floor is already at max timestamp.", c.element());
tracker.tryClaim(
@@ -962,8 +1012,9 @@ public ProcessContinuation process(
// Produce a poll result that only contains never seen before results in timestamp
// sorted order.
- Growth.PollResult newResults =
+ PendingResults pendingResults =
computeNeverSeenBeforeResults(pollingRestriction, res);
+ Growth.PollResult newResults = pendingResults.pollResult;
// If we had zero new results, attempt to update the watermark if the poll result
// provided a watermark or the retention floor bounds future outputs. Otherwise attempt
@@ -1015,20 +1066,30 @@ public ProcessContinuation process(
.max(
cursor,
newResults.getOutputs().get(newResults.getOutputs().size() - 1).getTimestamp());
- if (retentionFloorAtMaxTimestamp(newCursor, allowedLateness)) {
+ if (retentionFloorAtMaxTimestamp(newCursor, allowedLateness) && !pendingResults.truncated) {
LOG.info("{} - will stop polling, retention floor reached max timestamp.", c.element());
return stop();
}
}
Instant currentTime = Instant.now();
- if (getTerminationCondition().canStopPolling(currentTime, terminationState)) {
+ boolean terminationRequested =
+ getTerminationCondition().canStopPolling(currentTime, terminationState);
+ if (!pendingResults.truncated && terminationRequested) {
LOG.info(
"{} - told to stop polling by polling function at {} with termination state {}.",
c.element(),
currentTime,
getTerminationCondition().toString(terminationState));
return stop();
+ } else if (pendingResults.truncated && terminationRequested) {
+ LOG.info(
+ "{} - deferring termination until omitted outputs from the truncated poll are recovered.",
+ c.element());
+ }
+
+ if (pendingResults.truncated && BoundedWindow.TIMESTAMP_MAX_VALUE.equals(computedWatermark)) {
+ computedWatermark = BoundedWindow.TIMESTAMP_MAX_VALUE.minus(Duration.millis(1));
}
if (BoundedWindow.TIMESTAMP_MAX_VALUE.equals(computedWatermark)) {
@@ -1049,36 +1110,91 @@ private HashCode hash128(OutputT value) {
return Hashing.murmur3_128().hashObject(value, coderFunnel);
}
- private Growth.PollResult computeNeverSeenBeforeResults(
+ private PendingResults computeNeverSeenBeforeResults(
PollingGrowthState state, Growth.PollResult pollResult) {
// Collect results to include as newly pending. Note that the poll result may in theory
// contain multiple outputs mapping to the same output key - we need to ignore duplicates
// here already.
Instant retentionFloor = retentionFloor(state, spec.getTimestampCursorAllowedLateness());
- Map> newPending = Maps.newHashMap();
+ @Nullable Integer maxPendingResults = spec.getMaxPendingResults();
+ if (maxPendingResults == null
+ || BoundedWindow.TIMESTAMP_MAX_VALUE.equals(pollResult.getWatermark())) {
+ Map> newPending = Maps.newHashMap();
+ for (TimestampedValue output : pollResult.getOutputs()) {
+ if (retentionFloor != null && output.getTimestamp().isBefore(retentionFloor)) {
+ continue;
+ }
+ HashCode hash = hash128(output.getValue());
+ if (!state.getCompleted().containsKey(hash) && !newPending.containsKey(hash)) {
+ newPending.put(hash, output);
+ }
+ }
+ return new PendingResults<>(
+ pollResult.withOutputs(
+ Ordering.natural()
+ .onResultOf((TimestampedValue value) -> value.getTimestamp())
+ .sortedCopy(newPending.values())),
+ false);
+ }
+
+ Ordering> byTimestamp =
+ Ordering.natural().onResultOf((TimestampedValue value) -> value.getTimestamp());
+ PriorityQueue> oldestPending =
+ new PriorityQueue<>(
+ Math.min(maxPendingResults, Math.max(1, pollResult.getOutputs().size())),
+ byTimestamp.reverse());
+ Set hashesInPoll = new HashSet<>();
+ @Nullable Instant oldestOmittedTimestamp = null;
for (TimestampedValue output : pollResult.getOutputs()) {
if (retentionFloor != null && output.getTimestamp().isBefore(retentionFloor)) {
// The key that would prove this output already seen has been retired, so treat the
// output as seen.
continue;
}
- OutputT value = output.getValue();
- HashCode hash = hash128(value);
- if (state.getCompleted().containsKey(hash) || newPending.containsKey(hash)) {
+ HashCode hash = hash128(output.getValue());
+ if (state.getCompleted().containsKey(hash) || !hashesInPoll.add(hash)) {
continue;
}
- // TODO (https://github.com/apache/beam/issues/18459):
- // Consider adding only at most N pending elements and ignoring others,
- // instead relying on future poll rounds to provide them, in order to avoid
- // blowing up the state. Combined with a timestamp cursor, this would make the transform
- // scalable to very large poll results.
- newPending.put(hash, output);
- }
-
- return pollResult.withOutputs(
- Ordering.natural()
- .onResultOf((TimestampedValue value) -> value.getTimestamp())
- .sortedCopy(newPending.values()));
+ if (oldestPending.size() < maxPendingResults) {
+ oldestPending.add(output);
+ } else if (byTimestamp.compare(output, oldestPending.peek()) < 0) {
+ TimestampedValue omitted = oldestPending.remove();
+ if (oldestOmittedTimestamp == null
+ || omitted.getTimestamp().isBefore(oldestOmittedTimestamp)) {
+ oldestOmittedTimestamp = omitted.getTimestamp();
+ }
+ oldestPending.add(output);
+ } else {
+ if (oldestOmittedTimestamp == null
+ || output.getTimestamp().isBefore(oldestOmittedTimestamp)) {
+ oldestOmittedTimestamp = output.getTimestamp();
+ }
+ }
+ }
+
+ List> outputs = byTimestamp.sortedCopy(oldestPending);
+ if (oldestOmittedTimestamp == null) {
+ return new PendingResults<>(pollResult.withOutputs(outputs), false);
+ }
+
+ @Nullable Instant watermark = pollResult.getWatermark();
+ if (watermark != null && watermark.isAfter(oldestOmittedTimestamp)) {
+ watermark = oldestOmittedTimestamp;
+ }
+ if (BoundedWindow.TIMESTAMP_MAX_VALUE.equals(watermark)) {
+ watermark = watermark.minus(Duration.millis(1));
+ }
+ return new PendingResults<>(new PollResult<>(outputs, watermark), true);
+ }
+
+ private static final class PendingResults {
+ private final Growth.PollResult pollResult;
+ private final boolean truncated;
+
+ private PendingResults(Growth.PollResult pollResult, boolean truncated) {
+ this.pollResult = pollResult;
+ this.truncated = truncated;
+ }
}
private Growth.TerminationCondition getTerminationCondition() {
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/WatchTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/WatchTest.java
index ae32c6314e9e..51d050dc09b7 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/WatchTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/WatchTest.java
@@ -38,6 +38,7 @@
import java.io.IOException;
import java.io.Serializable;
import java.util.Arrays;
+import java.util.Collections;
import java.util.List;
import java.util.UUID;
import org.apache.beam.sdk.coders.Coder;
@@ -339,6 +340,35 @@ public PollResult> apply(String element, Context c)
p.run();
}
+ @Test
+ @Category({NeedsRunner.class, UsesUnboundedSplittableParDo.class})
+ public void testManyResultsWithPendingLimit() {
+ final int numResults = 300;
+ final Instant timestamp = Instant.now();
+ List all = Lists.newArrayList();
+ for (int i = 0; i < numResults; ++i) {
+ all.add(i);
+ }
+
+ PCollection> res =
+ p.apply(Create.of("a"))
+ .apply(
+ Watch.growthOf(
+ new PollFn() {
+ @Override
+ public PollResult apply(String element, Context c) {
+ return PollResult.incomplete(timestamp, all);
+ }
+ })
+ .withPollInterval(Duration.millis(1))
+ .withTerminationPerInput(Growth.afterIterations(4))
+ .withOutputCoder(VarIntCoder.of())
+ .withMaxPendingResults(100));
+
+ PAssert.that(res.apply(Values.create())).containsInAnyOrder(all);
+ p.run();
+ }
+
@Test
public void testCoder() throws Exception {
GrowthState pollingState =
@@ -612,10 +642,241 @@ public PollResult apply(String element, Context c) throws Exception {
new WatermarkEstimators.Manual(BoundedWindow.TIMESTAMP_MIN_VALUE);
ProcessContinuation processContinuation =
growthFn.process(context, tracker, watermarkEstimator);
+ verify(context)
+ .output(
+ KV.of(
+ null,
+ Arrays.asList(
+ TimestampedValue.of("a", now.plus(standardSeconds(1))),
+ TimestampedValue.of("b", now.plus(standardSeconds(2))),
+ TimestampedValue.of("c", now.plus(standardSeconds(3))),
+ TimestampedValue.of("d", now.plus(standardSeconds(4))))));
assertEquals(now.plus(standardSeconds(1)), watermarkEstimator.currentWatermark());
assertTrue(processContinuation.shouldResume());
}
+ @Test
+ public void testPollingGrowthTrackerDefersTerminationUntilTruncatedPollDrains() throws Exception {
+ Instant now = Instant.now();
+ List> outputs =
+ Arrays.asList(
+ TimestampedValue.of("d", now.plus(standardSeconds(4))),
+ TimestampedValue.of("c", now.plus(standardSeconds(3))),
+ TimestampedValue.of("a", now.plus(standardSeconds(1))),
+ TimestampedValue.of("b", now.plus(standardSeconds(2))));
+ Watch.Growth growth =
+ Watch.growthOf(
+ new PollFn() {
+ @Override
+ public PollResult apply(String element, Context c) {
+ return PollResult.incomplete(outputs)
+ .withWatermark(now.plus(standardSeconds(10)));
+ }
+ })
+ .withPollInterval(standardSeconds(10))
+ .withTerminationPerInput(Growth.afterIterations(1))
+ .withTimestampCursor()
+ .withMaxPendingResults(2);
+ WatchGrowthFn growthFn =
+ new WatchGrowthFn(
+ growth, StringUtf8Coder.of(), SerializableFunctions.identity(), StringUtf8Coder.of());
+
+ GrowthTracker firstTracker = newPollingGrowthTracker(Duration.ZERO);
+ DoFn.ProcessContext firstContext = mock(DoFn.ProcessContext.class);
+ ManualWatermarkEstimator firstWatermarkEstimator =
+ new WatermarkEstimators.Manual(BoundedWindow.TIMESTAMP_MIN_VALUE);
+
+ ProcessContinuation firstContinuation =
+ growthFn.process(firstContext, firstTracker, firstWatermarkEstimator);
+
+ verify(firstContext)
+ .output(
+ KV.of(
+ null,
+ Arrays.asList(
+ TimestampedValue.of("a", now.plus(standardSeconds(1))),
+ TimestampedValue.of("b", now.plus(standardSeconds(2))))));
+ assertEquals(now.plus(standardSeconds(3)), firstWatermarkEstimator.currentWatermark());
+ assertTrue(firstContinuation.shouldResume());
+
+ PollingGrowthState residual =
+ (PollingGrowthState) firstTracker.trySplit(0).getResidual();
+ assertEquals(
+ 2,
+ ((NonPollingGrowthState) firstTracker.currentRestriction())
+ .getPending()
+ .getOutputs()
+ .size());
+ GrowthTracker secondTracker = newTracker(residual, Duration.ZERO);
+ DoFn.ProcessContext secondContext = mock(DoFn.ProcessContext.class);
+ ManualWatermarkEstimator secondWatermarkEstimator =
+ new WatermarkEstimators.Manual(firstWatermarkEstimator.currentWatermark());
+
+ ProcessContinuation secondContinuation =
+ growthFn.process(secondContext, secondTracker, secondWatermarkEstimator);
+
+ verify(secondContext)
+ .output(
+ KV.of(
+ null,
+ Arrays.asList(
+ TimestampedValue.of("c", now.plus(standardSeconds(3))),
+ TimestampedValue.of("d", now.plus(standardSeconds(4))))));
+ assertEquals(
+ firstWatermarkEstimator.currentWatermark(), secondWatermarkEstimator.currentWatermark());
+ assertFalse(secondContinuation.shouldResume());
+ }
+
+ @Test
+ public void testPollingGrowthTrackerDoesNotLimitCompletePoll() throws Exception {
+ Instant now = Instant.now();
+ List> outputs =
+ Arrays.asList(
+ TimestampedValue.of("d", now.plus(standardSeconds(4))),
+ TimestampedValue.of("c", now.plus(standardSeconds(3))),
+ TimestampedValue.of("a", now.plus(standardSeconds(1))),
+ TimestampedValue.of("b", now.plus(standardSeconds(2))));
+ Watch.Growth growth =
+ Watch.growthOf(
+ new PollFn() {
+ @Override
+ public PollResult apply(String element, Context c) {
+ return PollResult.complete(outputs);
+ }
+ })
+ .withPollInterval(standardSeconds(10))
+ .withMaxPendingResults(2);
+ WatchGrowthFn growthFn =
+ new WatchGrowthFn(
+ growth, StringUtf8Coder.of(), SerializableFunctions.identity(), StringUtf8Coder.of());
+ DoFn.ProcessContext context = mock(DoFn.ProcessContext.class);
+ ManualWatermarkEstimator watermarkEstimator =
+ new WatermarkEstimators.Manual(BoundedWindow.TIMESTAMP_MIN_VALUE);
+
+ ProcessContinuation continuation =
+ growthFn.process(context, newPollingGrowthTracker(), watermarkEstimator);
+
+ verify(context)
+ .output(
+ KV.of(
+ null,
+ Arrays.asList(
+ TimestampedValue.of("a", now.plus(standardSeconds(1))),
+ TimestampedValue.of("b", now.plus(standardSeconds(2))),
+ TimestampedValue.of("c", now.plus(standardSeconds(3))),
+ TimestampedValue.of("d", now.plus(standardSeconds(4))))));
+ assertFalse(continuation.shouldResume());
+ }
+
+ @Test
+ public void testPollingGrowthTrackerPendingLimitKeepsFirstOutputForDuplicateKey()
+ throws Exception {
+ Instant now = Instant.now();
+ Watch.Growth growth =
+ Watch.growthOf(
+ new PollFn() {
+ @Override
+ public PollResult apply(String element, Context c) {
+ return PollResult.incomplete(
+ Arrays.asList(
+ TimestampedValue.of("a", now.plus(standardSeconds(4))),
+ TimestampedValue.of("a", now.plus(standardSeconds(1))),
+ TimestampedValue.of("b", now.plus(standardSeconds(2))),
+ TimestampedValue.of("c", now.plus(standardSeconds(3)))));
+ }
+ })
+ .withPollInterval(standardSeconds(10))
+ .withMaxPendingResults(2);
+ WatchGrowthFn growthFn =
+ new WatchGrowthFn(
+ growth, StringUtf8Coder.of(), SerializableFunctions.identity(), StringUtf8Coder.of());
+ DoFn.ProcessContext context = mock(DoFn.ProcessContext.class);
+ ManualWatermarkEstimator watermarkEstimator =
+ new WatermarkEstimators.Manual(BoundedWindow.TIMESTAMP_MIN_VALUE);
+
+ assertTrue(
+ growthFn.process(context, newPollingGrowthTracker(), watermarkEstimator).shouldResume());
+
+ verify(context)
+ .output(
+ KV.of(
+ null,
+ Arrays.asList(
+ TimestampedValue.of("b", now.plus(standardSeconds(2))),
+ TimestampedValue.of("c", now.plus(standardSeconds(3))))));
+ assertEquals(now.plus(standardSeconds(2)), watermarkEstimator.currentWatermark());
+ }
+
+ @Test
+ public void testPollingGrowthTrackerBoundsResultsAtMaxTimestamp() throws Exception {
+ List> outputs =
+ Arrays.asList(
+ TimestampedValue.of("a", BoundedWindow.TIMESTAMP_MAX_VALUE),
+ TimestampedValue.of("b", BoundedWindow.TIMESTAMP_MAX_VALUE));
+ Watch.Growth growth =
+ Watch.growthOf(
+ new PollFn() {
+ @Override
+ public PollResult apply(String element, Context c) {
+ return PollResult.incomplete(outputs);
+ }
+ })
+ .withPollInterval(standardSeconds(10))
+ .withTimestampCursor()
+ .withMaxPendingResults(1);
+ WatchGrowthFn growthFn =
+ new WatchGrowthFn(
+ growth, StringUtf8Coder.of(), SerializableFunctions.identity(), StringUtf8Coder.of());
+
+ GrowthTracker firstTracker = newPollingGrowthTracker(Duration.ZERO);
+ DoFn.ProcessContext firstContext = mock(DoFn.ProcessContext.class);
+ ManualWatermarkEstimator firstWatermarkEstimator =
+ new WatermarkEstimators.Manual(BoundedWindow.TIMESTAMP_MIN_VALUE);
+ assertTrue(
+ growthFn.process(firstContext, firstTracker, firstWatermarkEstimator).shouldResume());
+ verify(firstContext)
+ .output(
+ KV.of(
+ null, Arrays.asList(TimestampedValue.of("a", BoundedWindow.TIMESTAMP_MAX_VALUE))));
+ assertEquals(
+ BoundedWindow.TIMESTAMP_MAX_VALUE.minus(Duration.millis(1)),
+ firstWatermarkEstimator.currentWatermark());
+
+ PollingGrowthState residual =
+ (PollingGrowthState) firstTracker.trySplit(0).getResidual();
+ GrowthTracker secondTracker = newTracker(residual, Duration.ZERO);
+ DoFn.ProcessContext secondContext = mock(DoFn.ProcessContext.class);
+ ManualWatermarkEstimator secondWatermarkEstimator =
+ new WatermarkEstimators.Manual(firstWatermarkEstimator.currentWatermark());
+ assertFalse(
+ growthFn.process(secondContext, secondTracker, secondWatermarkEstimator).shouldResume());
+ verify(secondContext)
+ .output(
+ KV.of(
+ null, Arrays.asList(TimestampedValue.of("b", BoundedWindow.TIMESTAMP_MAX_VALUE))));
+ }
+
+ @Test
+ public void testMaxPendingResultsConfiguration() {
+ Watch.Growth growth =
+ Watch.growthOf(
+ new PollFn() {
+ @Override
+ public PollResult apply(String element, Context c) {
+ return PollResult.incomplete(Collections.emptyList());
+ }
+ });
+
+ assertNull(growth.getMaxPendingResults());
+ assertNull(growth.withMaxPendingResults(null).getMaxPendingResults());
+ Growth limited = growth.withMaxPendingResults(2);
+ assertEquals(Integer.valueOf(2), limited.getMaxPendingResults());
+ assertNull(limited.withMaxPendingResults(null).getMaxPendingResults());
+ assertEquals(Integer.valueOf(2), limited.getMaxPendingResults());
+ assertThrows(IllegalArgumentException.class, () -> growth.withMaxPendingResults(0));
+ assertThrows(IllegalArgumentException.class, () -> growth.withMaxPendingResults(-1));
+ }
+
@Test
public void testPollingGrowthTrackerDropsOutputsBehindCursor() throws Exception {
Instant now = Instant.now();