You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by yh...@apache.org on 2025/07/28 14:11:47 UTC
(beam) branch master updated: Revert "Refactor: separate SplittableTruncateSizedRestrictions" (#35700)
This is an automated email from the ASF dual-hosted git repository.
yhu 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 800f78bf5e9 Revert "Refactor: separate SplittableTruncateSizedRestrictions" (#35700)
800f78bf5e9 is described below
commit 800f78bf5e95dbaba36f8eef59cfde7391c8c944
Author: Yi Hu <ya...@google.com>
AuthorDate: Mon Jul 28 10:11:40 2025 -0400
Revert "Refactor: separate SplittableTruncateSizedRestrictions" (#35700)
---
.../sdk/fn/splittabledofn/RestrictionTrackers.java | 21 -
.../splittabledofn/RestrictionTracker.java | 2 +-
.../main/java/org/apache/beam/sdk/util/Holder.java | 45 -
.../main/java/org/apache/beam/sdk/values/KV.java | 1 -
.../org/apache/beam/sdk/values/WindowedValue.java | 7 -
.../apache/beam/fn/harness/FnApiDoFnRunner.java | 602 ++++++-
.../org/apache/beam/fn/harness/HandlesSplits.java | 2 -
.../org/apache/beam/fn/harness/ProgressUtils.java | 43 -
.../beam/fn/harness/SplitResultsWithStopIndex.java | 8 +-
...littableSplitAndSizeRestrictionsDoFnRunner.java | 2 +-
...ittableTruncateSizedRestrictionsDoFnRunner.java | 968 -----------
.../beam/fn/harness/WindowedSplitResult.java | 13 +-
.../beam/fn/harness/FnApiDoFnRunnerTest.java | 1744 ++++++++++++++++----
.../apache/beam/fn/harness/ProgressUtilsTest.java | 54 -
...bleTruncateSizedRestrictionsDoFnRunnerTest.java | 1191 -------------
15 files changed, 1990 insertions(+), 2713 deletions(-)
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/splittabledofn/RestrictionTrackers.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/splittabledofn/RestrictionTrackers.java
index 8879392d42a..6a80ce5ae3f 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/splittabledofn/RestrictionTrackers.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/splittabledofn/RestrictionTrackers.java
@@ -119,25 +119,4 @@ public class RestrictionTrackers {
return new RestrictionTrackerObserver<>(restrictionTracker, claimObserver);
}
}
-
- public static <RestrictionT, PositionT> RestrictionTracker<RestrictionT, PositionT> synchronize(
- RestrictionTracker<RestrictionT, PositionT> restrictionTracker) {
- if (restrictionTracker instanceof RestrictionTracker.HasProgress) {
- return new RestrictionTrackerObserverWithProgress<>(
- restrictionTracker, (ClaimObserver<PositionT>) NOOP_CLAIM_OBSERVER);
- } else {
- return new RestrictionTrackerObserver<>(
- restrictionTracker, (ClaimObserver<PositionT>) NOOP_CLAIM_OBSERVER);
- }
- }
-
- static class NoopClaimObserver<PositionT> implements ClaimObserver<PositionT> {
- @Override
- public void onClaimed(PositionT position) {}
-
- @Override
- public void onClaimFailed(PositionT position) {}
- }
-
- private static final NoopClaimObserver<Object> NOOP_CLAIM_OBSERVER = new NoopClaimObserver<>();
}
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/splittabledofn/RestrictionTracker.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/splittabledofn/RestrictionTracker.java
index 0c0b9f72647..8ab8d3ac8b4 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/splittabledofn/RestrictionTracker.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/splittabledofn/RestrictionTracker.java
@@ -211,6 +211,6 @@ public abstract class RestrictionTracker<RestrictionT, PositionT> {
return new AutoValue_RestrictionTracker_TruncateResult(restriction);
}
- public abstract RestrictionT getTruncatedRestriction();
+ public abstract @Nullable RestrictionT getTruncatedRestriction();
}
}
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/Holder.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/Holder.java
deleted file mode 100644
index 8a2853c6874..00000000000
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/Holder.java
+++ /dev/null
@@ -1,45 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.sdk.util;
-
-import org.apache.beam.sdk.annotations.Internal;
-
-/**
- * A trivial boxing of a value, used when nullability needs to be added to a generic type. (Optional
- * does not work for this)
- *
- * <p>Example: For a generic type `T` the actual parameter may be nullable or not. So you cannot
- * check values for null to determine presence/absence. Instead you can store a {@code @Nullable
- * Holder<T>}.
- */
-@Internal
-public class Holder<T> {
- private T value;
-
- private Holder(T value) {
- this.value = value;
- }
-
- public static <ValueT> Holder<ValueT> of(ValueT value) {
- return new Holder<>(value);
- }
-
- public T get() {
- return value;
- };
-}
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/KV.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/KV.java
index da32ae3b0e9..0744c2e008e 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/KV.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/KV.java
@@ -40,7 +40,6 @@ import org.checkerframework.dataflow.qual.Pure;
*/
public class KV<K, V> implements Serializable {
/** Returns a {@link KV} with the given key and value. */
- @Pure
public static <K, V> KV<K, V> of(K key, V value) {
return new KV<>(key, value);
}
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValue.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValue.java
index d2c4d1f07da..1dfa5feb7fd 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValue.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValue.java
@@ -20,7 +20,6 @@ package org.apache.beam.sdk.values;
import java.util.Collection;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.PaneInfo;
-import org.checkerframework.dataflow.qual.Pure;
import org.joda.time.Instant;
/**
@@ -30,32 +29,26 @@ import org.joda.time.Instant;
*/
public interface WindowedValue<T> {
/** The primary data for this value. */
- @Pure
T getValue();
/** The timestamp of this value in event time. */
- @Pure
Instant getTimestamp();
/** Returns the windows of this {@code WindowedValue}. */
- @Pure
Collection<? extends BoundedWindow> getWindows();
/** The {@link PaneInfo} associated with this WindowedValue. */
- @Pure
PaneInfo getPaneInfo();
/**
* A representation of each of the actual values represented by this compressed {@link
* WindowedValue}, one per window.
*/
- @Pure
Iterable<WindowedValue<T>> explodeWindows();
/**
* A {@link WindowedValue} with identical metadata to the current one, but with the provided
* value.
*/
- @Pure
<OtherT> WindowedValue<OtherT> withValue(OtherT value);
}
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java
index cc7c971e10b..4258681a494 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java
@@ -90,6 +90,7 @@ import org.apache.beam.sdk.transforms.reflect.DoFnSignatures;
import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.Progress;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.TruncateResult;
import org.apache.beam.sdk.transforms.splittabledofn.SplitResult;
import org.apache.beam.sdk.transforms.splittabledofn.TimestampObservingWatermarkEstimator;
import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
@@ -145,6 +146,7 @@ public class FnApiDoFnRunner<InputT, RestrictionT, PositionT, WatermarkEstimator
Factory factory = new Factory();
return ImmutableMap.<String, PTransformRunnerFactory>builder()
.put(PTransformTranslation.PAR_DO_TRANSFORM_URN, factory)
+ .put(PTransformTranslation.SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN, factory)
.put(
PTransformTranslation.SPLITTABLE_PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS_URN, factory)
.build();
@@ -249,7 +251,8 @@ public class FnApiDoFnRunner<InputT, RestrictionT, PositionT, WatermarkEstimator
/**
* Only valid during {@link
- * #processElementForWindowObservingSizedElementAndRestriction(WindowedValue)}.
+ * #processElementForWindowObservingSizedElementAndRestriction(WindowedValue)} and {@link
+ * #processElementForWindowObservingTruncateRestriction(WindowedValue)}.
*/
private List<BoundedWindow> currentWindows;
@@ -258,7 +261,8 @@ public class FnApiDoFnRunner<InputT, RestrictionT, PositionT, WatermarkEstimator
* processed.
*
* <p>Only valid during {@link
- * #processElementForWindowObservingSizedElementAndRestriction(WindowedValue)}.
+ * #processElementForWindowObservingSizedElementAndRestriction(WindowedValue)} and {@link
+ * #processElementForWindowObservingTruncateRestriction(WindowedValue)}.
*/
private int windowStopIndex;
@@ -267,17 +271,28 @@ public class FnApiDoFnRunner<InputT, RestrictionT, PositionT, WatermarkEstimator
* windowStopIndex.
*
* <p>Only valid during {@link
- * #processElementForWindowObservingSizedElementAndRestriction(WindowedValue)}.
+ * #processElementForWindowObservingSizedElementAndRestriction(WindowedValue)} and {@link
+ * #processElementForWindowObservingTruncateRestriction(WindowedValue)}.
*/
private int windowCurrentIndex;
- /** Only valid during #processElementForWindowObservingSizedElementAndRestriction}. */
+ /**
+ * Only valid during {@link #processElementForPairWithRestriction}, {@link
+ * #processElementForSplitRestriction}, and {@link
+ * #processElementForWindowObservingSizedElementAndRestriction}, null otherwise.
+ */
private RestrictionT currentRestriction;
- /** Only valid during {@link #processElementForWindowObservingSizedElementAndRestriction}. */
+ /**
+ * Only valid during {@link #processElementForSplitRestriction}, and {@link
+ * #processElementForWindowObservingSizedElementAndRestriction}, null otherwise.
+ */
private WatermarkEstimatorStateT currentWatermarkEstimatorState;
- /** Only valid during {@link #processElementForWindowObservingSizedElementAndRestriction}. */
+ /**
+ * Only valid during {@link #processElementForWindowObservingSizedElementAndRestriction} and
+ * {@link #processElementForWindowObservingTruncateRestriction}.
+ */
private Instant initialWatermark;
/**
@@ -346,6 +361,7 @@ public class FnApiDoFnRunner<InputT, RestrictionT, PositionT, WatermarkEstimator
mainOutputTag = (TupleTag) ParDoTranslation.getMainOutputTag(parDoPayload);
break;
case PTransformTranslation.SPLITTABLE_SPLIT_AND_SIZE_RESTRICTIONS_URN:
+ case PTransformTranslation.SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN:
mainOutputTag =
new TupleTag(Iterables.getOnlyElement(pTransform.getOutputsMap().keySet()));
break;
@@ -447,6 +463,8 @@ public class FnApiDoFnRunner<InputT, RestrictionT, PositionT, WatermarkEstimator
break;
case PTransformTranslation.SPLITTABLE_SPLIT_AND_SIZE_RESTRICTIONS_URN:
// startBundle should not be invoked
+ case PTransformTranslation.SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN:
+ // startBundle should not be invoked
default:
// no-op
}
@@ -468,6 +486,79 @@ public class FnApiDoFnRunner<InputT, RestrictionT, PositionT, WatermarkEstimator
this.processContext = new NonWindowObservingProcessBundleContext();
}
break;
+ case PTransformTranslation.SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN:
+ if ((doFnSignature.truncateRestriction() != null
+ && doFnSignature.truncateRestriction().observesWindow())
+ || (doFnSignature.newTracker() != null && doFnSignature.newTracker().observesWindow())
+ || (doFnSignature.getSize() != null && doFnSignature.getSize().observesWindow())
+ || !sideInputMapping.isEmpty()) {
+ // Only forward split/progress when the only consumer is splittable.
+ if (mainOutputConsumer instanceof HandlesSplits) {
+ mainInputConsumer =
+ new SplittableFnDataReceiver() {
+ private final HandlesSplits splitDelegate = (HandlesSplits) mainOutputConsumer;
+
+ @Override
+ public void accept(WindowedValue input) throws Exception {
+ processElementForWindowObservingTruncateRestriction(input);
+ }
+
+ @Override
+ public HandlesSplits.SplitResult trySplit(double fractionOfRemainder) {
+ return trySplitForWindowObservingTruncateRestriction(
+ fractionOfRemainder, splitDelegate);
+ }
+
+ @Override
+ public double getProgress() {
+ Progress progress =
+ FnApiDoFnRunner.this.getProgressFromWindowObservingTruncate(
+ splitDelegate.getProgress());
+ if (progress != null) {
+ double totalWork = progress.getWorkCompleted() + progress.getWorkRemaining();
+ if (totalWork > 0) {
+ return progress.getWorkCompleted() / totalWork;
+ }
+ }
+ return 0;
+ }
+ };
+ } else {
+ mainInputConsumer = this::processElementForWindowObservingTruncateRestriction;
+ }
+ this.processContext =
+ new SizedRestrictionWindowObservingProcessBundleContext(
+ PTransformTranslation.SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN);
+ } else {
+ // Only forward split/progress when the only consumer is splittable.
+ if (mainOutputConsumer instanceof HandlesSplits) {
+ mainInputConsumer =
+ new SplittableFnDataReceiver() {
+ private final HandlesSplits splitDelegate = (HandlesSplits) mainOutputConsumer;
+
+ @Override
+ public void accept(WindowedValue input) throws Exception {
+ processElementForTruncateRestriction(input);
+ }
+
+ @Override
+ public HandlesSplits.SplitResult trySplit(double fractionOfRemainder) {
+ return splitDelegate.trySplit(fractionOfRemainder);
+ }
+
+ @Override
+ public double getProgress() {
+ return splitDelegate.getProgress();
+ }
+ };
+ } else {
+ mainInputConsumer = this::processElementForTruncateRestriction;
+ }
+ this.processContext =
+ new SizedRestrictionNonWindowObservingProcessBundleContext(
+ PTransformTranslation.SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN);
+ }
+ break;
case PTransformTranslation.SPLITTABLE_PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS_URN:
if (doFnSignature.processElement().observesWindow()
|| (doFnSignature.newTracker() != null && doFnSignature.newTracker().observesWindow())
@@ -509,6 +600,8 @@ public class FnApiDoFnRunner<InputT, RestrictionT, PositionT, WatermarkEstimator
break;
case PTransformTranslation.SPLITTABLE_SPLIT_AND_SIZE_RESTRICTIONS_URN:
// finishBundle should not be invoked
+ case PTransformTranslation.SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN:
+ // finishBundle should not be invoked
default:
// no-op
}
@@ -643,6 +736,90 @@ public class FnApiDoFnRunner<InputT, RestrictionT, PositionT, WatermarkEstimator
}
}
+ private void processElementForTruncateRestriction(
+ WindowedValue<KV<KV<InputT, KV<RestrictionT, WatermarkEstimatorStateT>>, Double>> elem) {
+ currentElement = elem.withValue(elem.getValue().getKey().getKey());
+ currentRestriction = elem.getValue().getKey().getValue().getKey();
+ currentWatermarkEstimatorState = elem.getValue().getKey().getValue().getValue();
+ // For truncation, we don't set currentTrackerClaimed so that we enable checkpointing even if no
+ // progress is made.
+ currentTracker =
+ RestrictionTrackers.observe(
+ doFnInvoker.invokeNewTracker(processContext),
+ new ClaimObserver<PositionT>() {
+ @Override
+ public void onClaimed(PositionT position) {}
+
+ @Override
+ public void onClaimFailed(PositionT position) {}
+ });
+ try {
+ TruncateResult<OutputT> truncatedRestriction =
+ doFnInvoker.invokeTruncateRestriction(processContext);
+ if (truncatedRestriction != null) {
+ processContext.output(truncatedRestriction.getTruncatedRestriction());
+ }
+ } finally {
+ currentTracker = null;
+ currentElement = null;
+ currentRestriction = null;
+ currentWatermarkEstimatorState = null;
+ }
+
+ this.stateAccessor.finalizeState();
+ }
+
+ private void processElementForWindowObservingTruncateRestriction(
+ WindowedValue<KV<KV<InputT, KV<RestrictionT, WatermarkEstimatorStateT>>, Double>> elem) {
+ currentElement = elem.withValue(elem.getValue().getKey().getKey());
+ windowCurrentIndex = -1;
+ windowStopIndex = currentElement.getWindows().size();
+ currentWindows = ImmutableList.copyOf(currentElement.getWindows());
+ while (true) {
+ synchronized (splitLock) {
+ windowCurrentIndex++;
+ if (windowCurrentIndex >= windowStopIndex) {
+ // Careful to reset the split state under the same synchronized block.
+ windowCurrentIndex = -1;
+ windowStopIndex = 0;
+ currentElement = null;
+ currentWindows = null;
+ currentRestriction = null;
+ currentWatermarkEstimatorState = null;
+ currentWindow = null;
+ currentTracker = null;
+ currentWatermarkEstimator = null;
+ initialWatermark = null;
+ break;
+ }
+ currentRestriction = elem.getValue().getKey().getValue().getKey();
+ currentWatermarkEstimatorState = elem.getValue().getKey().getValue().getValue();
+ currentWindow = currentWindows.get(windowCurrentIndex);
+ // We leave currentTrackerClaimed unset as we want to split regardless of if tryClaim is
+ // called.
+ currentTracker =
+ RestrictionTrackers.observe(
+ doFnInvoker.invokeNewTracker(processContext),
+ new ClaimObserver<PositionT>() {
+ @Override
+ public void onClaimed(PositionT position) {}
+
+ @Override
+ public void onClaimFailed(PositionT position) {}
+ });
+ currentWatermarkEstimator =
+ WatermarkEstimators.threadSafe(doFnInvoker.invokeNewWatermarkEstimator(processContext));
+ initialWatermark = currentWatermarkEstimator.getWatermarkAndState().getKey();
+ }
+ TruncateResult<OutputT> truncatedRestriction =
+ doFnInvoker.invokeTruncateRestriction(processContext);
+ if (truncatedRestriction != null) {
+ processContext.output(truncatedRestriction.getTruncatedRestriction());
+ }
+ }
+ this.stateAccessor.finalizeState();
+ }
+
private void processElementForWindowObservingSizedElementAndRestriction(
WindowedValue<KV<KV<InputT, KV<RestrictionT, WatermarkEstimatorStateT>>, Double>> elem) {
currentElement = elem.withValue(elem.getValue().getKey().getKey());
@@ -745,13 +922,41 @@ public class FnApiDoFnRunner<InputT, RestrictionT, PositionT, WatermarkEstimator
private Progress getProgress() {
synchronized (splitLock) {
if (currentTracker instanceof RestrictionTracker.HasProgress && currentWindow != null) {
- return ProgressUtils.scaleProgress(
+ return scaleProgress(
((HasProgress) currentTracker).getProgress(), windowCurrentIndex, windowStopIndex);
}
}
return null;
}
+ private Progress getProgressFromWindowObservingTruncate(double elementCompleted) {
+ synchronized (splitLock) {
+ if (currentWindow != null) {
+ return scaleProgress(
+ Progress.from(elementCompleted, 1 - elementCompleted),
+ windowCurrentIndex,
+ windowStopIndex);
+ }
+ }
+ return null;
+ }
+
+ @VisibleForTesting
+ static Progress scaleProgress(Progress progress, int currentWindowIndex, int stopWindowIndex) {
+ checkArgument(
+ currentWindowIndex < stopWindowIndex,
+ "Current window index (%s) must be less than stop window index (%s)",
+ currentWindowIndex,
+ stopWindowIndex);
+
+ double totalWorkPerWindow = progress.getWorkCompleted() + progress.getWorkRemaining();
+ double completed = totalWorkPerWindow * currentWindowIndex + progress.getWorkCompleted();
+ double remaining =
+ totalWorkPerWindow * (stopWindowIndex - currentWindowIndex - 1)
+ + progress.getWorkRemaining();
+ return Progress.from(completed, remaining);
+ }
+
private WindowedSplitResult calculateRestrictionSize(
WindowedSplitResult splitResult, String errorContext) {
double fullSize =
@@ -835,6 +1040,61 @@ public class FnApiDoFnRunner<InputT, RestrictionT, PositionT, WatermarkEstimator
splitResult.getResidualInUnprocessedWindowsRoot().getPaneInfo()));
}
+ private HandlesSplits.SplitResult trySplitForWindowObservingTruncateRestriction(
+ double fractionOfRemainder, HandlesSplits splitDelegate) {
+ WindowedSplitResult windowedSplitResult = null;
+ HandlesSplits.SplitResult downstreamSplitResult = null;
+ synchronized (splitLock) {
+ // There is nothing to split if we are between truncate processing calls.
+ if (currentWindow == null) {
+ return null;
+ }
+ // We are requesting a checkpoint but have not yet progressed on the restriction, skip
+ // request.
+ if (fractionOfRemainder == 0
+ && currentTrackerClaimed != null
+ && !currentTrackerClaimed.get()) {
+ return null;
+ }
+
+ SplitResultsWithStopIndex splitResult =
+ computeSplitForProcessOrTruncate(
+ currentElement,
+ currentRestriction,
+ currentWindow,
+ currentWindows,
+ currentWatermarkEstimatorState,
+ fractionOfRemainder,
+ null,
+ splitDelegate,
+ null,
+ windowCurrentIndex,
+ windowStopIndex);
+ if (splitResult == null) {
+ return null;
+ }
+ windowStopIndex = splitResult.getNewWindowStopIndex();
+ windowedSplitResult =
+ calculateRestrictionSize(
+ splitResult.getWindowSplit(),
+ PTransformTranslation.SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN + "/GetSize");
+ downstreamSplitResult = splitResult.getDownstreamSplit();
+ }
+ // Note that the assumption here is the fullInputCoder of the Truncate transform should be the
+ // the same as the SDF/Process transform.
+ Coder fullInputCoder = WindowedValues.getFullCoder(inputCoder, windowCoder);
+ return constructSplitResult(
+ windowedSplitResult,
+ downstreamSplitResult,
+ fullInputCoder,
+ initialWatermark,
+ null,
+ pTransformId,
+ mainInputId,
+ pTransform.getOutputsMap().keySet(),
+ null);
+ }
+
private static <WatermarkEstimatorStateT> WindowedSplitResult computeWindowSplitResult(
WindowedValue currentElement,
Object currentRestriction,
@@ -892,7 +1152,7 @@ public class FnApiDoFnRunner<InputT, RestrictionT, PositionT, WatermarkEstimator
}
@VisibleForTesting
- static <WatermarkEstimatorStateT> SplitResultsWithStopIndex computeSplitForProcess(
+ static <WatermarkEstimatorStateT> SplitResultsWithStopIndex computeSplitForProcessOrTruncate(
WindowedValue currentElement,
Object currentRestriction,
BoundedWindow currentWindow,
@@ -929,8 +1189,7 @@ public class FnApiDoFnRunner<InputT, RestrictionT, PositionT, WatermarkEstimator
double elementCompleted = splitDelegate.getProgress();
elementProgress = Progress.from(elementCompleted, 1 - elementCompleted);
}
- Progress scaledProgress =
- ProgressUtils.scaleProgress(elementProgress, currentWindowIndex, stopWindowIndex);
+ Progress scaledProgress = scaleProgress(elementProgress, currentWindowIndex, stopWindowIndex);
double scaledFractionOfRemainder = scaledProgress.getWorkRemaining() * fractionOfRemainder;
// The fraction is out of the current window and hence we will split at the closest window
@@ -1158,7 +1417,7 @@ public class FnApiDoFnRunner<InputT, RestrictionT, PositionT, WatermarkEstimator
// applies to the residual.
watermarkAndState = currentWatermarkEstimator.getWatermarkAndState();
SplitResultsWithStopIndex splitResult =
- computeSplitForProcess(
+ computeSplitForProcessOrTruncate(
currentElement,
currentRestriction,
currentWindow,
@@ -1842,6 +2101,327 @@ public class FnApiDoFnRunner<InputT, RestrictionT, PositionT, WatermarkEstimator
}
}
+ /** This context outputs KV<KV<Element, KV<Restriction, WatemarkEstimatorState>>, Size>. */
+ private class SizedRestrictionWindowObservingProcessBundleContext
+ extends WindowObservingProcessBundleContextBase {
+ private final String errorContextPrefix;
+
+ SizedRestrictionWindowObservingProcessBundleContext(String errorContextPrefix) {
+ this.errorContextPrefix = errorContextPrefix;
+ }
+
+ @Override
+ // OutputT == RestrictionT
+ public void output(OutputT output) {
+ double size =
+ doFnInvoker.invokeGetSize(
+ new DelegatingArgumentProvider<InputT, OutputT>(
+ this, this.errorContextPrefix + "/GetSize") {
+ @Override
+ public Object restriction() {
+ return output;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return currentElement.getTimestamp();
+ }
+
+ @Override
+ public RestrictionTracker<?, ?> restrictionTracker() {
+ return doFnInvoker.invokeNewTracker(this);
+ }
+ });
+
+ // Don't need to check timestamp since we can always output using the input timestamp.
+ outputTo(
+ mainOutputConsumer,
+ (WindowedValue<OutputT>)
+ WindowedValues.of(
+ KV.of(
+ KV.of(
+ currentElement.getValue(), KV.of(output, currentWatermarkEstimatorState)),
+ size),
+ currentElement.getTimestamp(),
+ currentWindow,
+ currentElement.getPaneInfo()));
+ }
+
+ @Override
+ public <T> void output(TupleTag<T> tag, T output) {
+ // Note that the OutputReceiver/RowOutputReceiver specifically will use the non-tag versions
+ // of these methods when producing output.
+ throw new UnsupportedOperationException(
+ String.format("Non-main output %s unsupported in %s", tag, errorContextPrefix));
+ }
+
+ @Override
+ // OutputT == RestrictionT
+ public void outputWithTimestamp(OutputT output, Instant timestamp) {
+ checkTimestamp(timestamp);
+ double size =
+ doFnInvoker.invokeGetSize(
+ new DelegatingArgumentProvider<InputT, OutputT>(
+ this, this.errorContextPrefix + "/GetSize") {
+ @Override
+ public Object restriction() {
+ return output;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public RestrictionTracker<?, ?> restrictionTracker() {
+ return doFnInvoker.invokeNewTracker(this);
+ }
+ });
+
+ outputTo(
+ mainOutputConsumer,
+ (WindowedValue<OutputT>)
+ WindowedValues.of(
+ KV.of(
+ KV.of(
+ currentElement.getValue(), KV.of(output, currentWatermarkEstimatorState)),
+ size),
+ timestamp,
+ currentWindow,
+ currentElement.getPaneInfo()));
+ }
+
+ @Override
+ public void outputWindowedValue(
+ OutputT output,
+ Instant timestamp,
+ Collection<? extends BoundedWindow> windows,
+ PaneInfo paneInfo) {
+ checkTimestamp(timestamp);
+ double size =
+ doFnInvoker.invokeGetSize(
+ new DelegatingArgumentProvider<InputT, OutputT>(
+ this, this.errorContextPrefix + "/GetSize") {
+ @Override
+ public Object restriction() {
+ return output;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public RestrictionTracker<?, ?> restrictionTracker() {
+ return doFnInvoker.invokeNewTracker(this);
+ }
+ });
+
+ outputTo(
+ mainOutputConsumer,
+ (WindowedValue<OutputT>)
+ WindowedValues.of(
+ KV.of(
+ KV.of(
+ currentElement.getValue(), KV.of(output, currentWatermarkEstimatorState)),
+ size),
+ timestamp,
+ windows,
+ paneInfo));
+ }
+
+ @Override
+ public <T> void outputWithTimestamp(TupleTag<T> tag, T output, Instant timestamp) {
+ // Note that the OutputReceiver/RowOutputReceiver specifically will use the non-tag versions
+ // of these methods when producing output.
+ throw new UnsupportedOperationException(
+ String.format("Non-main output %s unsupported in %s", tag, errorContextPrefix));
+ }
+
+ @Override
+ public <T> void outputWindowedValue(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ Collection<? extends BoundedWindow> windows,
+ PaneInfo paneInfo) {
+ // Note that the OutputReceiver/RowOutputReceiver specifically will use the non-tag versions
+ // of these methods when producing output.
+ throw new UnsupportedOperationException(
+ String.format("Non-main output %s unsupported in %s", tag, errorContextPrefix));
+ }
+
+ @Override
+ public State state(String stateId, boolean alwaysFetched) {
+ throw new UnsupportedOperationException(
+ String.format("State unsupported in %s", errorContextPrefix));
+ }
+
+ @Override
+ public org.apache.beam.sdk.state.Timer timer(String timerId) {
+ throw new UnsupportedOperationException(
+ String.format("Timer unsupported in %s", errorContextPrefix));
+ }
+
+ @Override
+ public TimerMap timerFamily(String tagId) {
+ throw new UnsupportedOperationException(
+ String.format("Timer unsupported in %s", errorContextPrefix));
+ }
+ }
+
+ /** This context outputs KV<KV<Element, KV<Restriction, WatermarkEstimatorState>>, Size>. */
+ private class SizedRestrictionNonWindowObservingProcessBundleContext
+ extends NonWindowObservingProcessBundleContextBase {
+ private final String errorContextPrefix;
+
+ SizedRestrictionNonWindowObservingProcessBundleContext(String errorContextPrefix) {
+ this.errorContextPrefix = errorContextPrefix;
+ }
+
+ @Override
+ // OutputT == RestrictionT
+ public void output(OutputT output) {
+ double size =
+ doFnInvoker.invokeGetSize(
+ new DelegatingArgumentProvider<InputT, OutputT>(
+ this, errorContextPrefix + "/GetSize") {
+ @Override
+ public Object restriction() {
+ return output;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return currentElement.getTimestamp();
+ }
+
+ @Override
+ public RestrictionTracker<?, ?> restrictionTracker() {
+ return doFnInvoker.invokeNewTracker(this);
+ }
+ });
+
+ // Don't need to check timestamp since we can always output using the input timestamp.
+ outputTo(
+ mainOutputConsumer,
+ (WindowedValue<OutputT>)
+ currentElement.withValue(
+ KV.of(
+ KV.of(
+ currentElement.getValue(), KV.of(output, currentWatermarkEstimatorState)),
+ size)));
+ }
+
+ @Override
+ public <T> void output(TupleTag<T> tag, T output) {
+ // Note that the OutputReceiver/RowOutputReceiver specifically will use the non-tag versions
+ // of these methods when producing output.
+ throw new UnsupportedOperationException(
+ String.format("Non-main output %s unsupported in %s", tag, errorContextPrefix));
+ }
+
+ @Override
+ // OutputT == RestrictionT
+ public void outputWithTimestamp(OutputT output, Instant timestamp) {
+ checkTimestamp(timestamp);
+ double size =
+ doFnInvoker.invokeGetSize(
+ new DelegatingArgumentProvider<InputT, OutputT>(
+ this, errorContextPrefix + "/GetSize") {
+ @Override
+ public Object restriction() {
+ return output;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public RestrictionTracker<?, ?> restrictionTracker() {
+ return doFnInvoker.invokeNewTracker(this);
+ }
+ });
+
+ outputTo(
+ mainOutputConsumer,
+ (WindowedValue<OutputT>)
+ WindowedValues.of(
+ KV.of(
+ KV.of(
+ currentElement.getValue(), KV.of(output, currentWatermarkEstimatorState)),
+ size),
+ timestamp,
+ currentElement.getWindows(),
+ currentElement.getPaneInfo()));
+ }
+
+ @Override
+ public void outputWindowedValue(
+ OutputT output,
+ Instant timestamp,
+ Collection<? extends BoundedWindow> windows,
+ PaneInfo paneInfo) {
+ checkTimestamp(timestamp);
+ double size =
+ doFnInvoker.invokeGetSize(
+ new DelegatingArgumentProvider<InputT, OutputT>(
+ this, errorContextPrefix + "/GetSize") {
+ @Override
+ public Object restriction() {
+ return output;
+ }
+
+ @Override
+ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+ return timestamp;
+ }
+
+ @Override
+ public RestrictionTracker<?, ?> restrictionTracker() {
+ return doFnInvoker.invokeNewTracker(this);
+ }
+ });
+
+ outputTo(
+ mainOutputConsumer,
+ (WindowedValue<OutputT>)
+ WindowedValues.of(
+ KV.of(
+ KV.of(
+ currentElement.getValue(), KV.of(output, currentWatermarkEstimatorState)),
+ size),
+ timestamp,
+ windows,
+ paneInfo));
+ }
+
+ @Override
+ public <T> void outputWithTimestamp(TupleTag<T> tag, T output, Instant timestamp) {
+ // Note that the OutputReceiver/RowOutputReceiver specifically will use the non-tag versions
+ // of these methods when producing output.
+ throw new UnsupportedOperationException(
+ String.format("Non-main output %s unsupported in %s", tag, errorContextPrefix));
+ }
+
+ @Override
+ public <T> void outputWindowedValue(
+ TupleTag<T> tag,
+ T output,
+ Instant timestamp,
+ Collection<? extends BoundedWindow> windows,
+ PaneInfo paneInfo) {
+ // Note that the OutputReceiver/RowOutputReceiver specifically will use the non-tag versions
+ // of these methods when producing output.
+ throw new UnsupportedOperationException(
+ String.format("Non-main output %s unsupported in %s", tag, errorContextPrefix));
+ }
+ }
+
/** Provides arguments for a {@link DoFnInvoker} for a non-window observing method. */
private class NonWindowObservingProcessBundleContext
extends NonWindowObservingProcessBundleContextBase {
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/HandlesSplits.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/HandlesSplits.java
index af7638d61ca..54e19837e01 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/HandlesSplits.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/HandlesSplits.java
@@ -22,7 +22,6 @@ import java.util.List;
import org.apache.beam.model.fnexecution.v1.BeamFnApi;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.BundleApplication;
import org.apache.beam.sdk.fn.data.FnDataReceiver;
-import org.checkerframework.checker.nullness.qual.Nullable;
/**
* An interface that may be used to extend a {@link FnDataReceiver} signalling that the downstream
@@ -31,7 +30,6 @@ import org.checkerframework.checker.nullness.qual.Nullable;
public interface HandlesSplits {
/** Returns null if the split was unsuccessful. */
- @Nullable
SplitResult trySplit(double fractionOfRemainder);
/** Returns the current progress of the active element as a fraction between 0.0 and 1.0. */
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/ProgressUtils.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/ProgressUtils.java
deleted file mode 100644
index 6895df746ad..00000000000
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/ProgressUtils.java
+++ /dev/null
@@ -1,43 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.fn.harness;
-
-import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
-
-import org.apache.beam.sdk.annotations.Internal;
-import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
-
-/** Miscellaneous methods for working with progress. */
-@Internal
-public abstract class ProgressUtils {
- static RestrictionTracker.Progress scaleProgress(
- RestrictionTracker.Progress progress, int currentWindowIndex, int stopWindowIndex) {
- checkArgument(
- currentWindowIndex < stopWindowIndex,
- "Current window index (%s) must be less than stop window index (%s)",
- currentWindowIndex,
- stopWindowIndex);
-
- double totalWorkPerWindow = progress.getWorkCompleted() + progress.getWorkRemaining();
- double completed = totalWorkPerWindow * currentWindowIndex + progress.getWorkCompleted();
- double remaining =
- totalWorkPerWindow * (stopWindowIndex - currentWindowIndex - 1)
- + progress.getWorkRemaining();
- return RestrictionTracker.Progress.from(completed, remaining);
- }
-}
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/SplitResultsWithStopIndex.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/SplitResultsWithStopIndex.java
index a686aa99286..b0d6b5eee10 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/SplitResultsWithStopIndex.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/SplitResultsWithStopIndex.java
@@ -20,7 +20,6 @@ package org.apache.beam.fn.harness;
import com.google.auto.value.AutoValue;
import org.apache.beam.sdk.annotations.Internal;
import org.checkerframework.checker.nullness.qual.Nullable;
-import org.checkerframework.dataflow.qual.Pure;
@AutoValue
@AutoValue.CopyAnnotations
@@ -28,18 +27,15 @@ import org.checkerframework.dataflow.qual.Pure;
abstract class SplitResultsWithStopIndex {
public static SplitResultsWithStopIndex of(
WindowedSplitResult windowSplit,
- HandlesSplits.@Nullable SplitResult downstreamSplit,
+ HandlesSplits.SplitResult downstreamSplit,
int newWindowStopIndex) {
return new AutoValue_SplitResultsWithStopIndex(
windowSplit, downstreamSplit, newWindowStopIndex);
}
- @Pure
- public abstract WindowedSplitResult getWindowSplit();
+ public abstract @Nullable WindowedSplitResult getWindowSplit();
- @Pure
public abstract HandlesSplits.@Nullable SplitResult getDownstreamSplit();
- @Pure
public abstract int getNewWindowStopIndex();
}
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/SplittableSplitAndSizeRestrictionsDoFnRunner.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/SplittableSplitAndSizeRestrictionsDoFnRunner.java
index 0fd03447f6e..b701c468abc 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/SplittableSplitAndSizeRestrictionsDoFnRunner.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/SplittableSplitAndSizeRestrictionsDoFnRunner.java
@@ -359,7 +359,7 @@ public class SplittableSplitAndSizeRestrictionsDoFnRunner<
/** This context outputs KV<KV<Element, KV<Restriction, WatermarkEstimatorState>>, Size>. */
private class SizedRestrictionNonWindowObservingArgumentProvider
- extends SplitRestrictionArgumentProvider {
+ extends SplitRestrictionArgumentProvider implements OutputReceiver<RestrictionT> {
@Override
public void output(RestrictionT subrestriction) {
double size = getSize(subrestriction);
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/SplittableTruncateSizedRestrictionsDoFnRunner.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/SplittableTruncateSizedRestrictionsDoFnRunner.java
deleted file mode 100644
index f1001118d33..00000000000
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/SplittableTruncateSizedRestrictionsDoFnRunner.java
+++ /dev/null
@@ -1,968 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.fn.harness;
-
-import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
-import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
-
-import com.google.auto.service.AutoService;
-import java.io.IOException;
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import org.apache.beam.fn.harness.state.FnApiStateAccessor;
-import org.apache.beam.model.fnexecution.v1.BeamFnApi.BundleApplication;
-import org.apache.beam.model.fnexecution.v1.BeamFnApi.DelayedBundleApplication;
-import org.apache.beam.model.pipeline.v1.RunnerApi;
-import org.apache.beam.model.pipeline.v1.RunnerApi.PTransform;
-import org.apache.beam.model.pipeline.v1.RunnerApi.ParDoPayload;
-import org.apache.beam.sdk.Pipeline;
-import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.sdk.fn.data.FnDataReceiver;
-import org.apache.beam.sdk.fn.splittabledofn.RestrictionTrackers;
-import org.apache.beam.sdk.fn.splittabledofn.WatermarkEstimators;
-import org.apache.beam.sdk.options.PipelineOptions;
-import org.apache.beam.sdk.transforms.DoFn;
-import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
-import org.apache.beam.sdk.transforms.DoFnSchemaInformation;
-import org.apache.beam.sdk.transforms.SerializableFunction;
-import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
-import org.apache.beam.sdk.transforms.reflect.DoFnInvoker.DelegatingArgumentProvider;
-import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
-import org.apache.beam.sdk.transforms.reflect.DoFnSignature;
-import org.apache.beam.sdk.transforms.reflect.DoFnSignatures;
-import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
-import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.Progress;
-import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.TruncateResult;
-import org.apache.beam.sdk.transforms.splittabledofn.TimestampObservingWatermarkEstimator;
-import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
-import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
-import org.apache.beam.sdk.transforms.windowing.PaneInfo;
-import org.apache.beam.sdk.util.ByteStringOutputStream;
-import org.apache.beam.sdk.util.Holder;
-import org.apache.beam.sdk.util.UserCodeException;
-import org.apache.beam.sdk.util.construction.PTransformTranslation;
-import org.apache.beam.sdk.util.construction.ParDoTranslation;
-import org.apache.beam.sdk.util.construction.RehydratedComponents;
-import org.apache.beam.sdk.values.KV;
-import org.apache.beam.sdk.values.PCollectionView;
-import org.apache.beam.sdk.values.TupleTag;
-import org.apache.beam.sdk.values.WindowedValue;
-import org.apache.beam.sdk.values.WindowedValues;
-import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
-import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
-import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
-import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
-import org.checkerframework.checker.nullness.qual.NonNull;
-import org.checkerframework.checker.nullness.qual.Nullable;
-import org.joda.time.Instant;
-
-/**
- * A runner for the PTransform that truncates sized restrictions, for the case of draining a
- * pipeline.
- *
- * <p>The input and output types for this transform are
- * <li>{@code WindowedValue<KV<KV<InputT, KV<RestrictionT, WatermarkEstimatorStateT>>, Double>>}
- */
-public class SplittableTruncateSizedRestrictionsDoFnRunner<
- InputT, RestrictionT extends @NonNull Object, PositionT, WatermarkEstimatorStateT, OutputT>
- implements FnApiStateAccessor.MutatingStateContext<Void, BoundedWindow> {
-
- /** A registrar which provides a factory to handle Java {@link DoFn}s. */
- @AutoService(PTransformRunnerFactory.Registrar.class)
- public static class Registrar implements PTransformRunnerFactory.Registrar {
- @Override
- public Map<String, PTransformRunnerFactory> getPTransformRunnerFactories() {
- Factory factory = new Factory();
- return ImmutableMap.<String, PTransformRunnerFactory>builder()
- .put(PTransformTranslation.SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN, factory)
- .build();
- }
- }
-
- static class Factory implements PTransformRunnerFactory {
- @Override
- public final void addRunnerForPTransform(Context context) throws IOException {
- addRunnerForTruncateSizedRestrictions(context);
- }
-
- private <
- InputT,
- RestrictionT extends @NonNull Object,
- PositionT,
- WatermarkEstimatorStateT,
- OutputT>
- void addRunnerForTruncateSizedRestrictions(Context context) throws IOException {
-
- FnApiStateAccessor<Void> stateAccessor =
- FnApiStateAccessor.Factory.<Void>factoryForPTransformContext(context).create();
-
- // Main output
- checkArgument(
- context.getPTransform().getOutputsMap().size() == 1,
- "TruncateSizedRestrictions expects exact one output, but got: ",
- context.getPTransform().getOutputsMap().size());
- TupleTag<KV<KV<InputT, KV<RestrictionT, WatermarkEstimatorStateT>>, Double>> mainOutputTag =
- new TupleTag<>(
- Iterables.getOnlyElement(context.getPTransform().getOutputsMap().keySet()));
-
- FnDataReceiver<
- WindowedValue<KV<KV<InputT, KV<RestrictionT, WatermarkEstimatorStateT>>, Double>>>
- mainOutputConsumer =
- context.getPCollectionConsumer(
- context.getPTransform().getOutputsOrThrow(mainOutputTag.getId()));
-
- SplittableTruncateSizedRestrictionsDoFnRunner<
- InputT, RestrictionT, PositionT, WatermarkEstimatorStateT, OutputT>
- runner =
- new SplittableTruncateSizedRestrictionsDoFnRunner<>(
- context.getPipelineOptions(),
- context.getPTransformId(),
- context.getPTransform(),
- context.getComponents(),
- mainOutputConsumer,
- stateAccessor);
-
- // Register input consumer that delegates splitting
- FnDataReceiver<
- WindowedValue<KV<KV<InputT, KV<RestrictionT, WatermarkEstimatorStateT>>, Double>>>
- mainInputConsumer;
-
- if (mainOutputConsumer instanceof HandlesSplits) {
- mainInputConsumer =
- new SplitDelegatingFnDataReceiver<>(runner, (HandlesSplits) mainOutputConsumer);
- } else {
- mainInputConsumer = runner::processElement;
- }
-
- context.addPCollectionConsumer(
- context
- .getPTransform()
- .getInputsOrThrow(ParDoTranslation.getMainInputName(context.getPTransform())),
- mainInputConsumer);
- context.addTearDownFunction(runner::tearDown);
- }
- }
-
- //////////////////////////////////////////////////////////////////////////////////////////////////
-
- private final boolean observesWindow;
- private final PipelineOptions pipelineOptions;
-
- private final DoFnInvoker<InputT, OutputT> doFnInvoker;
-
- private final FnDataReceiver<
- WindowedValue<KV<KV<InputT, KV<RestrictionT, WatermarkEstimatorStateT>>, Double>>>
- mainOutputConsumer;
-
- private final FnApiStateAccessor<?> stateAccessor;
-
- private final TruncateSizedRestrictionArgumentProvider mutableArgumentProvider;
-
- private final String pTransformId;
- private final RunnerApi.PTransform pTransform;
- private final String mainInputId;
- private final Coder<WindowedValue<?>> fullInputCoder;
-
- /**
- * Used to guarantee a consistent view of this {@link
- * SplittableTruncateSizedRestrictionsDoFnRunner} while setting up for {@link
- * DoFnInvoker#invokeProcessElement} since {@link #trySplit} may access internal {@link
- * SplittableTruncateSizedRestrictionsDoFnRunner} state concurrently.
- */
- // TODO: explicitly mark guarded fields with @GuardedBy
- private final Object splitLock = new Object();
-
- private final DoFnSchemaInformation doFnSchemaInformation;
- private final Map<String, PCollectionView<?>> sideInputMapping;
-
- ///
- // Mutating fields that change with the element and window being processed
- //
- private int windowCurrentIndex;
- private @Nullable List<BoundedWindow> currentWindows;
- private @Nullable RestrictionT currentRestriction;
- private @Nullable Holder<WatermarkEstimatorStateT> currentWatermarkEstimatorState;
- private @Nullable Instant initialWatermark;
- private WatermarkEstimators.@Nullable WatermarkAndStateObserver<WatermarkEstimatorStateT>
- currentWatermarkEstimator;
- private @Nullable BoundedWindow currentWindow;
- private @Nullable RestrictionTracker<RestrictionT, PositionT> currentTracker;
- private @Nullable WindowedValue<InputT> currentElement;
-
- /**
- * The window index at which processing should stop. The window with this index should not be
- * processed.
- */
- private int windowStopIndex;
-
- /**
- * The window index which is currently being processed. This should always be less than
- * windowStopIndex.
- */
- SplittableTruncateSizedRestrictionsDoFnRunner(
- PipelineOptions pipelineOptions,
- String pTransformId,
- PTransform pTransform,
- RunnerApi.Components components,
- FnDataReceiver<
- WindowedValue<KV<KV<InputT, KV<RestrictionT, WatermarkEstimatorStateT>>, Double>>>
- mainOutputConsumer,
- FnApiStateAccessor<Void> stateAccessor)
- throws IOException {
- this.pipelineOptions = pipelineOptions;
- this.stateAccessor = stateAccessor;
- this.pTransformId = pTransformId;
- this.pTransform = pTransform;
-
- ParDoPayload parDoPayload = ParDoPayload.parseFrom(pTransform.getSpec().getPayload());
-
- // DoFn and metadata
- DoFn<InputT, OutputT> doFn = (DoFn<InputT, OutputT>) ParDoTranslation.getDoFn(parDoPayload);
- DoFnSignature doFnSignature = DoFnSignatures.signatureForDoFn(doFn);
- this.doFnInvoker = DoFnInvokers.tryInvokeSetupFor(doFn, pipelineOptions);
- this.doFnSchemaInformation = ParDoTranslation.getSchemaInformation(parDoPayload);
-
- this.mainOutputConsumer = mainOutputConsumer;
-
- // Side inputs
- this.sideInputMapping = ParDoTranslation.getSideInputMapping(parDoPayload);
-
- // Register processing methods
- this.observesWindow =
- (doFnSignature.splitRestriction() != null
- && doFnSignature.splitRestriction().observesWindow())
- || (doFnSignature.newTracker() != null && doFnSignature.newTracker().observesWindow())
- || (doFnSignature.getSize() != null && doFnSignature.getSize().observesWindow())
- || !sideInputMapping.isEmpty();
-
- if (observesWindow) {
- this.mutableArgumentProvider = new TruncateSizedRestrictionWindowObservingArgumentProvider();
- } else {
- this.mutableArgumentProvider =
- new TruncateSizedRestrictionNonWindowObservingArgumentProvider();
- }
-
- // Main Input
- this.mainInputId = ParDoTranslation.getMainInputName(pTransform);
- RunnerApi.PCollection mainInput =
- components.getPcollectionsOrThrow(pTransform.getInputsOrThrow(mainInputId));
- RehydratedComponents rehydratedComponents =
- RehydratedComponents.forComponents(components).withPipeline(Pipeline.create());
- Coder<KV<KV<InputT, KV<RestrictionT, WatermarkEstimatorStateT>>, Double>> inputCoder =
- (Coder<KV<KV<InputT, KV<RestrictionT, WatermarkEstimatorStateT>>, Double>>)
- rehydratedComponents.getCoder(mainInput.getCoderId());
- Coder<BoundedWindow> windowCoder =
- (Coder<BoundedWindow>)
- rehydratedComponents
- .getWindowingStrategy(mainInput.getWindowingStrategyId())
- .getWindowFn()
- .windowCoder();
- this.fullInputCoder =
- (Coder<WindowedValue<?>>) (Object) WindowedValues.getFullCoder(inputCoder, windowCoder);
- }
-
- @Override
- public Void getCurrentKey() {
- return null;
- }
-
- @Override
- public BoundedWindow getCurrentWindow() {
- return checkStateNotNull(
- currentWindow, "Attempt to access window outside windowed element processing context.");
- }
-
- public List<BoundedWindow> getCurrentWindows() {
- return checkStateNotNull(
- currentWindows,
- "Attempt to access window collection outside windowed element processing context.");
- }
-
- public WindowedValue<InputT> getCurrentElement() {
- return checkStateNotNull(
- currentElement, "Attempt to access element outside element processing context.");
- }
-
- private RestrictionT getCurrentRestriction() {
- return checkStateNotNull(
- this.currentRestriction,
- "Attempt to access restriction outside element processing context.");
- }
-
- private RestrictionTracker<RestrictionT, ?> getCurrentTracker() {
- return checkStateNotNull(
- this.currentTracker,
- "Attempt to access restriction tracker state outside element processing context.");
- }
-
- private WatermarkEstimatorStateT getCurrentWatermarkEstimatorState() {
- checkStateNotNull(
- this.currentWatermarkEstimatorState,
- "Attempt to access watermark estimator state outside element processing context.");
- return this.currentWatermarkEstimatorState.get();
- }
-
- void processElement(
- WindowedValue<KV<KV<InputT, KV<RestrictionT, WatermarkEstimatorStateT>>, Double>> elem) {
- if (observesWindow) {
- processElementForWindowObservingTruncateRestriction(elem);
- } else {
- processElementForTruncateRestriction(elem);
- }
- }
-
- HandlesSplits.@Nullable SplitResult trySplit(
- double fractionOfRemainder, HandlesSplits splitDelegate) {
- if (observesWindow) {
- return trySplitForWindowObservingTruncateRestriction(fractionOfRemainder, splitDelegate);
- } else {
- return splitDelegate.trySplit(fractionOfRemainder);
- }
- }
-
- private void processElementForTruncateRestriction(
- WindowedValue<KV<KV<InputT, KV<RestrictionT, WatermarkEstimatorStateT>>, Double>> elem) {
- currentElement = elem.withValue(elem.getValue().getKey().getKey());
- currentRestriction = elem.getValue().getKey().getValue().getKey();
- currentWatermarkEstimatorState = Holder.of(elem.getValue().getKey().getValue().getValue());
- currentTracker =
- RestrictionTrackers.synchronize(doFnInvoker.invokeNewTracker(mutableArgumentProvider));
- try {
- TruncateResult<RestrictionT> truncatedRestriction =
- doFnInvoker.invokeTruncateRestriction(mutableArgumentProvider);
- if (truncatedRestriction != null) {
- mutableArgumentProvider.output(truncatedRestriction.getTruncatedRestriction());
- }
- } finally {
- currentTracker = null;
- currentElement = null;
- currentRestriction = null;
- currentWatermarkEstimatorState = null;
- }
-
- this.stateAccessor.finalizeState();
- }
-
- private void processElementForWindowObservingTruncateRestriction(
- WindowedValue<KV<KV<InputT, KV<RestrictionT, WatermarkEstimatorStateT>>, Double>> elem) {
- currentElement = elem.withValue(elem.getValue().getKey().getKey());
- windowCurrentIndex = -1;
- windowStopIndex = elem.getWindows().size();
- currentWindows = ImmutableList.copyOf(elem.getWindows());
- while (true) {
- synchronized (splitLock) {
- windowCurrentIndex++;
- if (windowCurrentIndex >= windowStopIndex) {
- // Careful to reset the split state under the same synchronized block.
- windowCurrentIndex = -1;
- windowStopIndex = 0;
- currentElement = null;
- currentWindows = null;
- currentRestriction = null;
- currentWatermarkEstimatorState = null;
- currentWindow = null;
- currentTracker = null;
- currentWatermarkEstimator = null;
- initialWatermark = null;
- break;
- }
- currentRestriction = elem.getValue().getKey().getValue().getKey();
- currentWatermarkEstimatorState = Holder.of(elem.getValue().getKey().getValue().getValue());
- currentWindow =
- checkStateNotNull(
- currentWindows,
- "internal error: currentWindows is null during element processing")
- .get(windowCurrentIndex);
- currentTracker =
- RestrictionTrackers.synchronize(doFnInvoker.invokeNewTracker(mutableArgumentProvider));
- currentWatermarkEstimator =
- WatermarkEstimators.threadSafe(
- doFnInvoker.invokeNewWatermarkEstimator(mutableArgumentProvider));
- initialWatermark = currentWatermarkEstimator.getWatermarkAndState().getKey();
- }
- TruncateResult<RestrictionT> truncatedRestriction =
- doFnInvoker.invokeTruncateRestriction(mutableArgumentProvider);
- if (truncatedRestriction != null) {
- mutableArgumentProvider.output(truncatedRestriction.getTruncatedRestriction());
- }
- }
- this.stateAccessor.finalizeState();
- }
-
- private @Nullable Progress getProgressFromWindowObservingTruncate(double elementCompleted) {
- synchronized (splitLock) {
- if (currentWindow != null) {
- return ProgressUtils.scaleProgress(
- Progress.from(elementCompleted, 1 - elementCompleted),
- windowCurrentIndex,
- windowStopIndex);
- }
- }
- return null;
- }
-
- private WindowedSplitResult calculateRestrictionSize(
- WindowedSplitResult splitResult, String errorContext) {
- double fullSize =
- splitResult.getResidualInUnprocessedWindowsRoot() == null
- && splitResult.getPrimaryInFullyProcessedWindowsRoot() == null
- ? 0
- : doFnInvoker.invokeGetSize(
- new DelegatingArgumentProvider<InputT, OutputT>(
- mutableArgumentProvider, errorContext) {
- @Override
- public Object restriction() {
- return getCurrentRestriction();
- }
-
- @Override
- public RestrictionTracker<?, ?> restrictionTracker() {
- return doFnInvoker.invokeNewTracker(this);
- }
- });
- double primarySize =
- splitResult.getPrimarySplitRoot() == null
- ? 0
- : doFnInvoker.invokeGetSize(
- new DelegatingArgumentProvider<InputT, OutputT>(
- mutableArgumentProvider, errorContext) {
- @Override
- public Object restriction() {
- WindowedValue<KV<InputT, KV<RestrictionT, WatermarkEstimatorStateT>>>
- splitRoot =
- (WindowedValue<KV<InputT, KV<RestrictionT, WatermarkEstimatorStateT>>>)
- splitResult.getPrimarySplitRoot();
-
- return splitRoot.getValue().getValue().getKey();
- }
-
- @Override
- public RestrictionTracker<?, ?> restrictionTracker() {
- return doFnInvoker.invokeNewTracker(this);
- }
- });
- double residualSize =
- splitResult.getResidualSplitRoot() == null
- ? 0
- : doFnInvoker.invokeGetSize(
- new DelegatingArgumentProvider<InputT, OutputT>(
- mutableArgumentProvider, errorContext) {
- @Override
- public Object restriction() {
- WindowedValue<KV<InputT, KV<RestrictionT, WatermarkEstimatorStateT>>>
- splitRoot =
- (WindowedValue<KV<InputT, KV<RestrictionT, WatermarkEstimatorStateT>>>)
- splitResult.getResidualSplitRoot();
-
- return splitRoot.getValue().getValue().getKey();
- }
-
- @Override
- public RestrictionTracker<?, ?> restrictionTracker() {
- return doFnInvoker.invokeNewTracker(this);
- }
- });
- return WindowedSplitResult.forRoots(
- splitResult.getPrimaryInFullyProcessedWindowsRoot() == null
- ? null
- : WindowedValues.of(
- KV.of(splitResult.getPrimaryInFullyProcessedWindowsRoot().getValue(), fullSize),
- splitResult.getPrimaryInFullyProcessedWindowsRoot().getTimestamp(),
- splitResult.getPrimaryInFullyProcessedWindowsRoot().getWindows(),
- splitResult.getPrimaryInFullyProcessedWindowsRoot().getPaneInfo()),
- splitResult.getPrimarySplitRoot() == null
- ? null
- : WindowedValues.of(
- KV.of(splitResult.getPrimarySplitRoot().getValue(), primarySize),
- splitResult.getPrimarySplitRoot().getTimestamp(),
- splitResult.getPrimarySplitRoot().getWindows(),
- splitResult.getPrimarySplitRoot().getPaneInfo()),
- splitResult.getResidualSplitRoot() == null
- ? null
- : WindowedValues.of(
- KV.of(splitResult.getResidualSplitRoot().getValue(), residualSize),
- splitResult.getResidualSplitRoot().getTimestamp(),
- splitResult.getResidualSplitRoot().getWindows(),
- splitResult.getResidualSplitRoot().getPaneInfo()),
- splitResult.getResidualInUnprocessedWindowsRoot() == null
- ? null
- : WindowedValues.of(
- KV.of(splitResult.getResidualInUnprocessedWindowsRoot().getValue(), fullSize),
- splitResult.getResidualInUnprocessedWindowsRoot().getTimestamp(),
- splitResult.getResidualInUnprocessedWindowsRoot().getWindows(),
- splitResult.getResidualInUnprocessedWindowsRoot().getPaneInfo()));
- }
-
- private HandlesSplits.@Nullable SplitResult trySplitForWindowObservingTruncateRestriction(
- double fractionOfRemainder, HandlesSplits splitDelegate) {
- WindowedSplitResult windowedSplitResult;
- HandlesSplits.SplitResult downstreamSplitResult;
- synchronized (splitLock) {
- // There is nothing to split if we are between truncate processing calls.
- if (currentWindow == null) {
- return null;
- }
-
- SplitResultsWithStopIndex splitResult =
- computeSplitForTruncate(
- getCurrentElement(),
- getCurrentRestriction(),
- getCurrentWindow(),
- getCurrentWindows(),
- getCurrentWatermarkEstimatorState(),
- fractionOfRemainder,
- splitDelegate,
- windowCurrentIndex,
- windowStopIndex);
- if (splitResult == null) {
- return null;
- }
- windowStopIndex = splitResult.getNewWindowStopIndex();
- windowedSplitResult =
- calculateRestrictionSize(
- splitResult.getWindowSplit(),
- PTransformTranslation.SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN + "/GetSize");
- downstreamSplitResult = splitResult.getDownstreamSplit();
- }
- return constructSplitResult(
- windowedSplitResult,
- downstreamSplitResult,
- fullInputCoder,
- checkStateNotNull(
- initialWatermark, "Attempt to construct split result without initial watermark"),
- pTransformId,
- mainInputId,
- pTransform.getOutputsMap().keySet());
- }
-
- private static <WatermarkEstimatorStateT> WindowedSplitResult computeWindowSplitResult(
- WindowedValue<?> currentElement,
- Object currentRestriction,
- List<BoundedWindow> windows,
- WatermarkEstimatorStateT currentWatermarkEstimatorState,
- int toIndex,
- int fromIndex,
- int stopWindowIndex) {
- List<BoundedWindow> primaryFullyProcessedWindows = windows.subList(0, toIndex);
- List<BoundedWindow> residualUnprocessedWindows = windows.subList(fromIndex, stopWindowIndex);
- WindowedSplitResult windowedSplitResult;
-
- windowedSplitResult =
- WindowedSplitResult.forRoots(
- primaryFullyProcessedWindows.isEmpty()
- ? null
- : WindowedValues.of(
- KV.of(
- currentElement.getValue(),
- KV.of(currentRestriction, currentWatermarkEstimatorState)),
- currentElement.getTimestamp(),
- primaryFullyProcessedWindows,
- currentElement.getPaneInfo()),
- null,
- null,
- residualUnprocessedWindows.isEmpty()
- ? null
- : WindowedValues.of(
- KV.of(
- currentElement.getValue(),
- KV.of(currentRestriction, currentWatermarkEstimatorState)),
- currentElement.getTimestamp(),
- residualUnprocessedWindows,
- currentElement.getPaneInfo()));
- return windowedSplitResult;
- }
-
- @VisibleForTesting
- static <WatermarkEstimatorStateT> @Nullable SplitResultsWithStopIndex computeSplitForTruncate(
- WindowedValue<?> element,
- Object restriction,
- BoundedWindow window,
- List<BoundedWindow> windows,
- WatermarkEstimatorStateT watermarkEstimatorState,
- double fractionOfRemainder,
- HandlesSplits splitDelegate,
- int currentWindowIndex,
- int stopWindowIndex) {
- checkArgument(splitDelegate != null);
-
- WindowedSplitResult windowedSplitResult;
- HandlesSplits.@Nullable SplitResult downstreamSplitResult = null;
- int newWindowStopIndex;
- // If we are not on the last window, try to compute the split which is on the current window or
- // on a future window.
- if (currentWindowIndex != stopWindowIndex - 1) {
- // Compute the fraction of the remainder relative to the scaled progress.
- Progress elementProgress;
- double elementCompleted = splitDelegate.getProgress();
- elementProgress = Progress.from(elementCompleted, 1 - elementCompleted);
- Progress scaledProgress =
- ProgressUtils.scaleProgress(elementProgress, currentWindowIndex, stopWindowIndex);
- double scaledFractionOfRemainder = scaledProgress.getWorkRemaining() * fractionOfRemainder;
-
- // The fraction is out of the current window and hence we will split at the closest window
- // boundary.
- if (scaledFractionOfRemainder >= elementProgress.getWorkRemaining()) {
- newWindowStopIndex =
- (int)
- Math.min(
- stopWindowIndex - 1,
- currentWindowIndex
- + Math.max(
- 1,
- Math.round(
- (elementProgress.getWorkCompleted() + scaledFractionOfRemainder)
- / (elementProgress.getWorkCompleted()
- + elementProgress.getWorkRemaining()))));
- windowedSplitResult =
- computeWindowSplitResult(
- element,
- restriction,
- windows,
- watermarkEstimatorState,
- newWindowStopIndex,
- newWindowStopIndex,
- stopWindowIndex);
- } else {
- // Compute the element split with the scaled fraction.
- downstreamSplitResult = splitDelegate.trySplit(scaledFractionOfRemainder);
- newWindowStopIndex = currentWindowIndex + 1;
- int toIndex = (downstreamSplitResult == null) ? newWindowStopIndex : currentWindowIndex;
- windowedSplitResult =
- computeWindowSplitResult(
- element,
- restriction,
- windows,
- watermarkEstimatorState,
- toIndex,
- newWindowStopIndex,
- stopWindowIndex);
- }
- } else {
- // We are on the last window then compute the element split with given fraction.
- newWindowStopIndex = stopWindowIndex;
- downstreamSplitResult = splitDelegate.trySplit(fractionOfRemainder);
- if (downstreamSplitResult == null) {
- return null;
- }
- windowedSplitResult =
- computeWindowSplitResult(
- element,
- restriction,
- windows,
- watermarkEstimatorState,
- currentWindowIndex,
- stopWindowIndex,
- stopWindowIndex);
- }
- return SplitResultsWithStopIndex.of(
- windowedSplitResult, downstreamSplitResult, newWindowStopIndex);
- }
-
- @VisibleForTesting
- static <WatermarkEstimatorStateT> HandlesSplits.SplitResult constructSplitResult(
- @Nullable WindowedSplitResult windowedSplitResult,
- HandlesSplits.@Nullable SplitResult downstreamElementSplit,
- Coder<WindowedValue<?>> fullInputCoder,
- Instant initialWatermark,
- String pTransformId,
- String mainInputId,
- Collection<String> outputIds) {
- // The element split cannot from both windowedSplitResult and downstreamElementSplit.
- checkArgument(
- (windowedSplitResult == null || windowedSplitResult.getResidualSplitRoot() == null)
- || downstreamElementSplit == null);
- List<BundleApplication> primaryRoots = new ArrayList<>();
- List<DelayedBundleApplication> residualRoots = new ArrayList<>();
-
- // Encode window splits.
- if (windowedSplitResult != null
- && windowedSplitResult.getPrimaryInFullyProcessedWindowsRoot() != null) {
- ByteStringOutputStream primaryInOtherWindowsBytes = new ByteStringOutputStream();
- try {
- fullInputCoder.encode(
- windowedSplitResult.getPrimaryInFullyProcessedWindowsRoot(),
- primaryInOtherWindowsBytes);
- } catch (IOException e) {
- throw new RuntimeException(e);
- }
- BundleApplication.Builder primaryApplicationInOtherWindows =
- BundleApplication.newBuilder()
- .setTransformId(pTransformId)
- .setInputId(mainInputId)
- .setElement(primaryInOtherWindowsBytes.toByteString());
- primaryRoots.add(primaryApplicationInOtherWindows.build());
- }
- if (windowedSplitResult != null
- && windowedSplitResult.getResidualInUnprocessedWindowsRoot() != null) {
- ByteStringOutputStream bytesOut = new ByteStringOutputStream();
- try {
- fullInputCoder.encode(windowedSplitResult.getResidualInUnprocessedWindowsRoot(), bytesOut);
- } catch (IOException e) {
- throw new RuntimeException(e);
- }
- BundleApplication.Builder residualInUnprocessedWindowsRoot =
- BundleApplication.newBuilder()
- .setTransformId(pTransformId)
- .setInputId(mainInputId)
- .setElement(bytesOut.toByteString());
- // We don't want to change the output watermarks or set the checkpoint resume time since
- // that applies to the current window.
- Map<String, org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.Timestamp>
- outputWatermarkMapForUnprocessedWindows = new HashMap<>();
- if (!initialWatermark.equals(GlobalWindow.TIMESTAMP_MIN_VALUE)) {
- org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.Timestamp outputWatermark =
- org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.Timestamp.newBuilder()
- .setSeconds(initialWatermark.getMillis() / 1000)
- .setNanos((int) (initialWatermark.getMillis() % 1000) * 1000000)
- .build();
- for (String outputId : outputIds) {
- outputWatermarkMapForUnprocessedWindows.put(outputId, outputWatermark);
- }
- }
- residualInUnprocessedWindowsRoot.putAllOutputWatermarks(
- outputWatermarkMapForUnprocessedWindows);
- residualRoots.add(
- DelayedBundleApplication.newBuilder()
- .setApplication(residualInUnprocessedWindowsRoot)
- .build());
- }
-
- // Encode element split from windowedSplitResult or from downstream element split. It's possible
- // that there is no element split.
- if (downstreamElementSplit != null) {
- primaryRoots.add(Iterables.getOnlyElement(downstreamElementSplit.getPrimaryRoots()));
- residualRoots.add(Iterables.getOnlyElement(downstreamElementSplit.getResidualRoots()));
- }
-
- return HandlesSplits.SplitResult.of(primaryRoots, residualRoots);
- }
-
- /** Outputs the given element to the specified set of consumers wrapping any exceptions. */
- private <T> void outputTo(FnDataReceiver<WindowedValue<T>> consumer, WindowedValue<T> output) {
- if (currentWatermarkEstimator instanceof TimestampObservingWatermarkEstimator) {
- ((TimestampObservingWatermarkEstimator) currentWatermarkEstimator)
- .observeTimestamp(output.getTimestamp());
- }
- try {
- consumer.accept(output);
- } catch (Throwable t) {
- throw UserCodeException.wrap(t);
- }
- }
-
- private void tearDown() {
- doFnInvoker.invokeTeardown();
- }
-
- /** This context outputs KV<KV<Element, KV<Restriction, WatemarkEstimatorState>>, Size>. */
- private class TruncateSizedRestrictionWindowObservingArgumentProvider
- extends TruncateSizedRestrictionArgumentProvider {
-
- @Override
- public void output(RestrictionT output) {
- double size = getSize(output);
- outputTo(
- mainOutputConsumer,
- WindowedValues.of(
- KV.of(
- KV.of(
- getCurrentElement().getValue(),
- KV.of(output, getCurrentWatermarkEstimatorState())),
- size),
- getCurrentElement().getTimestamp(),
- getCurrentWindow(),
- getCurrentElement().getPaneInfo()));
- }
-
- @Override
- public BoundedWindow window() {
- return getCurrentWindow();
- }
-
- @Override
- public Object sideInput(String tagId) {
- PCollectionView<Object> pCollectionView =
- (PCollectionView<Object>)
- checkStateNotNull(sideInputMapping.get(tagId), "Side input tag not found: %s", tagId);
-
- return stateAccessor.get(pCollectionView, getCurrentWindow());
- }
- }
-
- /** This context outputs KV<KV<Element, KV<Restriction, WatermarkEstimatorState>>, Size>. */
- private class TruncateSizedRestrictionNonWindowObservingArgumentProvider
- extends TruncateSizedRestrictionArgumentProvider {
-
- @Override
- public void output(RestrictionT truncatedRestriction) {
- double size = getSize(truncatedRestriction);
- outputTo(
- mainOutputConsumer,
- getCurrentElement()
- .withValue(
- KV.of(
- KV.of(
- getCurrentElement().getValue(),
- KV.of(truncatedRestriction, getCurrentWatermarkEstimatorState())),
- size)));
- }
- }
-
- /** Base implementation that does not override methods which need to be window aware. */
- private abstract class TruncateSizedRestrictionArgumentProvider
- extends DoFnInvoker.BaseArgumentProvider<InputT, OutputT>
- implements OutputReceiver<RestrictionT> {
-
- protected double getSize(RestrictionT subrestriction) {
- return doFnInvoker.invokeGetSize(
- new DelegatingArgumentProvider<InputT, OutputT>(this, getErrorContext() + "/GetSize") {
- @Override
- public Object restriction() {
- return subrestriction;
- }
-
- @Override
- public Instant timestamp(DoFn<InputT, OutputT> doFn) {
- return getCurrentElement().getTimestamp();
- }
-
- @Override
- public RestrictionTracker<?, ?> restrictionTracker() {
- return doFnInvoker.invokeNewTracker(this);
- }
- });
- }
-
- @Override
- public String getErrorContext() {
- return "TruncateRestriction";
- }
-
- @Override
- public PaneInfo paneInfo(DoFn<InputT, OutputT> doFn) {
- return getCurrentElement().getPaneInfo();
- }
-
- @Override
- public InputT element(DoFn<InputT, OutputT> doFn) {
- return getCurrentElement().getValue();
- }
-
- @Override
- public Object schemaElement(int index) {
- SerializableFunction<InputT, Object> converter =
- (SerializableFunction<InputT, Object>)
- doFnSchemaInformation.getElementConverters().get(index);
- return converter.apply(getCurrentElement().getValue());
- }
-
- @Override
- public Instant timestamp(DoFn<InputT, OutputT> doFn) {
- return getCurrentElement().getTimestamp();
- }
-
- @Override
- public OutputReceiver<OutputT> outputReceiver(DoFn<InputT, OutputT> doFn) {
- // OutputT == RestrictionT
- return (OutputReceiver<OutputT>) this;
- }
-
- @Override
- public Object restriction() {
- return getCurrentRestriction();
- }
-
- @Override
- @SuppressWarnings("nullness")
- public Object watermarkEstimatorState() {
- return getCurrentWatermarkEstimatorState();
- }
-
- @Override
- public RestrictionTracker<?, ?> restrictionTracker() {
- return getCurrentTracker();
- }
-
- @Override
- public PipelineOptions pipelineOptions() {
- return pipelineOptions;
- }
-
- @Override
- public void outputWithTimestamp(RestrictionT output, Instant timestamp) {
- throw new UnsupportedOperationException(
- "Cannot outputWithTimestamp from TruncateRestriction");
- }
- }
-
- /**
- * Passes split requests downstream, by way of trySplit on the overall runner, which will split
- * per window if needed.
- */
- private static class SplitDelegatingFnDataReceiver<
- InputT, RestrictionT extends @NonNull Object, WatermarkEstimatorStateT>
- implements FnDataReceiver<
- WindowedValue<KV<KV<InputT, KV<RestrictionT, WatermarkEstimatorStateT>>, Double>>>,
- HandlesSplits {
-
- private final HandlesSplits splitDelegate;
- private final SplittableTruncateSizedRestrictionsDoFnRunner<
- InputT, RestrictionT, ?, WatermarkEstimatorStateT, ?>
- runner;
-
- public SplitDelegatingFnDataReceiver(
- SplittableTruncateSizedRestrictionsDoFnRunner<
- InputT, RestrictionT, ?, WatermarkEstimatorStateT, ?>
- runner,
- HandlesSplits splitDelegate) {
- this.runner = runner;
- this.splitDelegate = splitDelegate;
- }
-
- @Override
- public void accept(
- WindowedValue<KV<KV<InputT, KV<RestrictionT, WatermarkEstimatorStateT>>, Double>> input) {
- runner.processElement(input);
- }
-
- @Override
- public @Nullable SplitResult trySplit(double fractionOfRemainder) {
- return runner.trySplit(fractionOfRemainder, splitDelegate);
- }
-
- @Override
- public double getProgress() {
- double delegateProgress = splitDelegate.getProgress();
- if (!runner.observesWindow) {
- return delegateProgress;
- } else {
- Progress progress = runner.getProgressFromWindowObservingTruncate(delegateProgress);
- if (progress != null) {
- double totalWork = progress.getWorkCompleted() + progress.getWorkRemaining();
- if (totalWork > 0) {
- return progress.getWorkCompleted() / totalWork;
- }
- }
- return 0;
- }
- }
- }
-}
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/WindowedSplitResult.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/WindowedSplitResult.java
index 82517781388..bf9eefd03a8 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/WindowedSplitResult.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/WindowedSplitResult.java
@@ -21,7 +21,6 @@ import com.google.auto.value.AutoValue;
import org.apache.beam.sdk.annotations.Internal;
import org.apache.beam.sdk.values.WindowedValue;
import org.checkerframework.checker.nullness.qual.Nullable;
-import org.checkerframework.dataflow.qual.Pure;
/** Internal class to hold the primary and residual roots when converted to an input element. */
@AutoValue
@@ -29,10 +28,10 @@ import org.checkerframework.dataflow.qual.Pure;
@Internal
abstract class WindowedSplitResult {
public static WindowedSplitResult forRoots(
- @Nullable WindowedValue<?> primaryInFullyProcessedWindowsRoot,
- @Nullable WindowedValue<?> primarySplitRoot,
- @Nullable WindowedValue<?> residualSplitRoot,
- @Nullable WindowedValue<?> residualInUnprocessedWindowsRoot) {
+ WindowedValue<?> primaryInFullyProcessedWindowsRoot,
+ WindowedValue<?> primarySplitRoot,
+ WindowedValue<?> residualSplitRoot,
+ WindowedValue<?> residualInUnprocessedWindowsRoot) {
return new AutoValue_WindowedSplitResult(
primaryInFullyProcessedWindowsRoot,
primarySplitRoot,
@@ -40,15 +39,11 @@ abstract class WindowedSplitResult {
residualInUnprocessedWindowsRoot);
}
- @Pure
public abstract @Nullable WindowedValue<?> getPrimaryInFullyProcessedWindowsRoot();
- @Pure
public abstract @Nullable WindowedValue<?> getPrimarySplitRoot();
- @Pure
public abstract @Nullable WindowedValue<?> getResidualSplitRoot();
- @Pure
public abstract @Nullable WindowedValue<?> getResidualInUnprocessedWindowsRoot();
}
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java
index edd9c465464..bdc21cdadfe 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java
@@ -103,6 +103,7 @@ import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.View;
import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.Progress;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.FixedWindows;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
@@ -2602,64 +2603,227 @@ public class FnApiDoFnRunnerTest implements Serializable {
}
}
- /**
- * A {@link DoFn} that outputs elements with timestamp equal to the input timestamp minus the
- * input element.
- */
- private static class SkewingDoFn extends DoFn<String, String> {
- private final Duration allowedSkew;
+ @Test
+ public void testProcessElementForTruncateAndSizeRestrictionForwardSplitWhenObservingWindows()
+ throws Exception {
+ Pipeline p = Pipeline.create();
+ PCollection<String> valuePCollection = p.apply(Create.of("unused"));
+ PCollectionView<String> singletonSideInputView = valuePCollection.apply(View.asSingleton());
+ WindowObservingTestSplittableDoFn doFn =
+ WindowObservingTestSplittableDoFn.forSplitAtTruncate(singletonSideInputView);
+ valuePCollection
+ .apply(Window.into(SlidingWindows.of(Duration.standardSeconds(1))))
+ .apply(TEST_TRANSFORM_ID, ParDo.of(doFn).withSideInputs(singletonSideInputView));
- private SkewingDoFn(Duration allowedSkew) {
- this.allowedSkew = allowedSkew;
- }
+ RunnerApi.Pipeline pProto =
+ ProtoOverrides.updateTransform(
+ PTransformTranslation.PAR_DO_TRANSFORM_URN,
+ PipelineTranslation.toProto(p, SdkComponents.create(p.getOptions()), true),
+ SplittableParDoExpander.createTruncateReplacement());
+ String expandedTransformId =
+ Iterables.find(
+ pProto.getComponents().getTransformsMap().entrySet(),
+ entry ->
+ entry
+ .getValue()
+ .getSpec()
+ .getUrn()
+ .equals(
+ PTransformTranslation.SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN)
+ && entry.getValue().getUniqueName().contains(TEST_TRANSFORM_ID))
+ .getKey();
+ RunnerApi.PTransform pTransform =
+ pProto.getComponents().getTransformsOrThrow(expandedTransformId);
+ String inputPCollectionId =
+ pTransform.getInputsOrThrow(ParDoTranslation.getMainInputName(pTransform));
+ RunnerApi.PCollection inputPCollection =
+ pProto.getComponents().getPcollectionsOrThrow(inputPCollectionId);
+ RehydratedComponents rehydratedComponents =
+ RehydratedComponents.forComponents(pProto.getComponents());
+ Coder<WindowedValue> inputCoder =
+ WindowedValues.getFullCoder(
+ CoderTranslation.fromProto(
+ pProto.getComponents().getCodersOrThrow(inputPCollection.getCoderId()),
+ rehydratedComponents,
+ TranslationContext.DEFAULT),
+ (Coder)
+ CoderTranslation.fromProto(
+ pProto
+ .getComponents()
+ .getCodersOrThrow(
+ pProto
+ .getComponents()
+ .getWindowingStrategiesOrThrow(
+ inputPCollection.getWindowingStrategyId())
+ .getWindowCoderId()),
+ rehydratedComponents,
+ TranslationContext.DEFAULT));
- @ProcessElement
- public void processElement(ProcessContext context) {
- Duration duration = Duration.millis(Long.valueOf(context.element()));
- context.outputWithTimestamp(context.element(), context.timestamp().minus(duration));
- }
+ String outputPCollectionId = Iterables.getOnlyElement(pTransform.getOutputsMap().values());
- @Override
- public Duration getAllowedTimestampSkew() {
- return allowedSkew;
- }
- }
+ FakeBeamFnStateClient fakeClient =
+ new FakeBeamFnStateClient(StringUtf8Coder.of(), ImmutableMap.of());
- private static class OutputFnDataReceiver implements FnDataReceiver<WindowedValue> {
- OutputFnDataReceiver(List<WindowedValue<String>> mainOutputValues) {
- this.mainOutputValues = mainOutputValues;
- }
+ PTransformRunnerFactoryTestContext context =
+ PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
+ .beamFnStateClient(fakeClient)
+ .processBundleInstructionId("57")
+ // .pCollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .putAllEnvironments(Collections.emptyMap())
+ .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .build())
+ // .coders(pProto.getComponents().getCodersMap())
+ // .windowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .build();
+ List<WindowedValue<KV<KV<String, OffsetRange>, Double>>> mainOutputValues = new ArrayList<>();
+ Coder coder =
+ KvCoder.of(
+ KvCoder.of(
+ StringUtf8Coder.of(), KvCoder.of(OffsetRange.Coder.of(), InstantCoder.of())),
+ DoubleCoder.of());
+ context.addPCollectionConsumer(
+ outputPCollectionId, (FnDataReceiver) new SplittableFnDataReceiver(mainOutputValues));
- private final List<WindowedValue<String>> mainOutputValues;
+ new FnApiDoFnRunner.Factory<>().addRunnerForPTransform(context);
+ FnDataReceiver<WindowedValue<?>> mainInput =
+ context.getPCollectionConsumer(inputPCollectionId);
+ assertThat(mainInput, instanceOf(HandlesSplits.class));
- @Override
- public void accept(WindowedValue input) throws Exception {
- mainOutputValues.add(input);
- }
+ mainOutputValues.clear();
+ BoundedWindow window1 = new IntervalWindow(new Instant(5), new Instant(10));
+ BoundedWindow window2 = new IntervalWindow(new Instant(6), new Instant(11));
+ BoundedWindow window3 = new IntervalWindow(new Instant(7), new Instant(12));
+ // Setup and launch the trySplit thread.
+ ExecutorService executorService = Executors.newSingleThreadExecutor();
+ Future<HandlesSplits.SplitResult> trySplitFuture =
+ executorService.submit(
+ () -> {
+ try {
+ doFn.waitForSplitElementToBeProcessed();
+ HandlesSplits.SplitResult result = ((HandlesSplits) mainInput).trySplit(0);
+ Assert.assertNotNull(result);
+ return result;
+ } finally {
+ doFn.trySplitPerformed();
+ }
+ });
+
+ WindowedValue<?> splitValue =
+ valueInWindows(
+ KV.of(
+ KV.of("7", KV.of(new OffsetRange(0, 6), GlobalWindow.TIMESTAMP_MIN_VALUE)), 6.0),
+ window1,
+ window2,
+ window3);
+ mainInput.accept(splitValue);
+ HandlesSplits.SplitResult trySplitResult = trySplitFuture.get();
+
+ // We expect that there are outputs from window1 and window2
+ assertThat(
+ mainOutputValues,
+ contains(
+ WindowedValues.of(
+ KV.of(
+ KV.of("7", KV.of(new OffsetRange(0, 3), GlobalWindow.TIMESTAMP_MIN_VALUE)),
+ 3.0),
+ splitValue.getTimestamp(),
+ window1,
+ splitValue.getPaneInfo()),
+ WindowedValues.of(
+ KV.of(
+ KV.of("7", KV.of(new OffsetRange(0, 3), GlobalWindow.TIMESTAMP_MIN_VALUE)),
+ 3.0),
+ splitValue.getTimestamp(),
+ window2,
+ splitValue.getPaneInfo())));
+
+ SplitResult expectedElementSplit = createSplitResult(0);
+ BundleApplication expectedElementSplitPrimary =
+ Iterables.getOnlyElement(expectedElementSplit.getPrimaryRoots());
+ ByteStringOutputStream primaryBytes = new ByteStringOutputStream();
+ inputCoder.encode(
+ WindowedValues.of(
+ KV.of(
+ KV.of("7", KV.of(new OffsetRange(0, 6), GlobalWindow.TIMESTAMP_MIN_VALUE)), 6.0),
+ splitValue.getTimestamp(),
+ window1,
+ splitValue.getPaneInfo()),
+ primaryBytes);
+ BundleApplication expectedWindowedPrimary =
+ BundleApplication.newBuilder()
+ .setElement(primaryBytes.toByteString())
+ .setInputId(ParDoTranslation.getMainInputName(pTransform))
+ .setTransformId(TEST_TRANSFORM_ID)
+ .build();
+ DelayedBundleApplication expectedElementSplitResidual =
+ Iterables.getOnlyElement(expectedElementSplit.getResidualRoots());
+ ByteStringOutputStream residualBytes = new ByteStringOutputStream();
+ inputCoder.encode(
+ WindowedValues.of(
+ KV.of(
+ KV.of("7", KV.of(new OffsetRange(0, 6), GlobalWindow.TIMESTAMP_MIN_VALUE)), 6.0),
+ splitValue.getTimestamp(),
+ window3,
+ splitValue.getPaneInfo()),
+ residualBytes);
+ DelayedBundleApplication expectedWindowedResidual =
+ DelayedBundleApplication.newBuilder()
+ .setApplication(
+ BundleApplication.newBuilder()
+ .setElement(residualBytes.toByteString())
+ .setInputId(ParDoTranslation.getMainInputName(pTransform))
+ .setTransformId(TEST_TRANSFORM_ID)
+ .build())
+ .build();
+ assertThat(
+ trySplitResult.getPrimaryRoots(),
+ contains(expectedWindowedPrimary, expectedElementSplitPrimary));
+ assertThat(
+ trySplitResult.getResidualRoots(),
+ contains(expectedWindowedResidual, expectedElementSplitResidual));
}
@Test
- public void testDoFnSkewNotAllowed() throws Exception {
+ public void testProcessElementForTruncateAndSizeRestrictionForwardSplitWithoutObservingWindow()
+ throws Exception {
Pipeline p = Pipeline.create();
- PCollection<String> valuePCollection = p.apply(Create.of("0", "1"));
- PCollection<String> outputPCollection =
- valuePCollection.apply(TEST_TRANSFORM_ID, ParDo.of(new SkewingDoFn(Duration.ZERO)));
+ PCollection<String> valuePCollection = p.apply(Create.of("unused"));
+ valuePCollection.apply(
+ TEST_TRANSFORM_ID, ParDo.of(new NonWindowObservingTestSplittableDoFn()));
- SdkComponents sdkComponents = SdkComponents.create(p.getOptions());
- RunnerApi.Pipeline pProto = PipelineTranslation.toProto(p, sdkComponents);
- String inputPCollectionId = sdkComponents.registerPCollection(valuePCollection);
- String outputPCollectionId = sdkComponents.registerPCollection(outputPCollection);
+ RunnerApi.Pipeline pProto =
+ ProtoOverrides.updateTransform(
+ PTransformTranslation.PAR_DO_TRANSFORM_URN,
+ PipelineTranslation.toProto(p, SdkComponents.create(p.getOptions()), true),
+ SplittableParDoExpander.createTruncateReplacement());
+ String expandedTransformId =
+ Iterables.find(
+ pProto.getComponents().getTransformsMap().entrySet(),
+ entry ->
+ entry
+ .getValue()
+ .getSpec()
+ .getUrn()
+ .equals(
+ PTransformTranslation.SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN)
+ && entry.getValue().getUniqueName().contains(TEST_TRANSFORM_ID))
+ .getKey();
RunnerApi.PTransform pTransform =
- pProto
- .getComponents()
- .getTransformsOrThrow(
- pProto
- .getComponents()
- .getTransformsOrThrow(TEST_TRANSFORM_ID)
- .getSubtransforms(0));
+ pProto.getComponents().getTransformsOrThrow(expandedTransformId);
+ String inputPCollectionId =
+ pTransform.getInputsOrThrow(ParDoTranslation.getMainInputName(pTransform));
+ String outputPCollectionId = Iterables.getOnlyElement(pTransform.getOutputsMap().values());
+
+ FakeBeamFnStateClient fakeClient =
+ new FakeBeamFnStateClient(StringUtf8Coder.of(), ImmutableMap.of());
PTransformRunnerFactoryTestContext context =
PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
+ .beamFnStateClient(fakeClient)
.processBundleInstructionId("57")
.components(
RunnerApi.Components.newBuilder()
@@ -2669,59 +2833,57 @@ public class FnApiDoFnRunnerTest implements Serializable {
.putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
.build())
.build();
- List<WindowedValue<String>> mainOutputValues = new ArrayList<>();
+ List<WindowedValue<KV<KV<String, OffsetRange>, Double>>> mainOutputValues = new ArrayList<>();
+ Coder coder =
+ KvCoder.of(KvCoder.of(StringUtf8Coder.of(), OffsetRange.Coder.of()), DoubleCoder.of());
context.addPCollectionConsumer(
- outputPCollectionId, (FnDataReceiver) new OutputFnDataReceiver(mainOutputValues));
+ outputPCollectionId, (FnDataReceiver) new SplittableFnDataReceiver(mainOutputValues));
new FnApiDoFnRunner.Factory<>().addRunnerForPTransform(context);
-
- mainOutputValues.clear();
FnDataReceiver<WindowedValue<?>> mainInput =
context.getPCollectionConsumer(inputPCollectionId);
- mainInput.accept(valueInGlobalWindow("0"));
-
- String message =
- assertThrows(
- UserCodeException.class,
- () -> {
- mainInput.accept(timestampedValueInGlobalWindow("1", new Instant(0L)));
- })
- .getMessage();
+ assertThat(mainInput, instanceOf(HandlesSplits.class));
- assertThat(
- message,
- allOf(
- containsString(
- String.format("timestamp %s", new Instant(0).minus(Duration.millis(1L)))),
- containsString(
- String.format(
- "allowed skew (%s)",
- PeriodFormat.getDefault().print(Duration.ZERO.toPeriod())))));
+ assertEquals(0.7, ((HandlesSplits) mainInput).getProgress(), 0.0);
+ assertEquals(createSplitResult(0.4), ((HandlesSplits) mainInput).trySplit(0.4));
}
@Test
- public void testDoFnSkewAllowed() throws Exception {
+ public void testProcessElementForTruncateAndSizeRestriction() throws Exception {
Pipeline p = Pipeline.create();
- PCollection<String> valuePCollection = p.apply(Create.of("0", "3"));
- PCollection<String> outputPCollection =
- valuePCollection.apply(TEST_TRANSFORM_ID, ParDo.of(new SkewingDoFn(Duration.millis(5L))));
+ PCollection<String> valuePCollection = p.apply(Create.of("unused"));
+ valuePCollection.apply(
+ TEST_TRANSFORM_ID, ParDo.of(new NonWindowObservingTestSplittableDoFn()));
- SdkComponents sdkComponents = SdkComponents.create(p.getOptions());
- RunnerApi.Pipeline pProto = PipelineTranslation.toProto(p, sdkComponents);
- String inputPCollectionId = sdkComponents.registerPCollection(valuePCollection);
- String outputPCollectionId = sdkComponents.registerPCollection(outputPCollection);
+ RunnerApi.Pipeline pProto =
+ ProtoOverrides.updateTransform(
+ PTransformTranslation.PAR_DO_TRANSFORM_URN,
+ PipelineTranslation.toProto(p, SdkComponents.create(p.getOptions()), true),
+ SplittableParDoExpander.createTruncateReplacement());
+ String expandedTransformId =
+ Iterables.find(
+ pProto.getComponents().getTransformsMap().entrySet(),
+ entry ->
+ entry
+ .getValue()
+ .getSpec()
+ .getUrn()
+ .equals(
+ PTransformTranslation.SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN)
+ && entry.getValue().getUniqueName().contains(TEST_TRANSFORM_ID))
+ .getKey();
RunnerApi.PTransform pTransform =
- pProto
- .getComponents()
- .getTransformsOrThrow(
- pProto
- .getComponents()
- .getTransformsOrThrow(TEST_TRANSFORM_ID)
- .getSubtransforms(0));
+ pProto.getComponents().getTransformsOrThrow(expandedTransformId);
+ String inputPCollectionId =
+ pTransform.getInputsOrThrow(ParDoTranslation.getMainInputName(pTransform));
+ String outputPCollectionId = Iterables.getOnlyElement(pTransform.getOutputsMap().values());
+
+ FakeBeamFnStateClient fakeClient =
+ new FakeBeamFnStateClient(StringUtf8Coder.of(), ImmutableMap.of());
- List<WindowedValue<String>> mainOutputValues = new ArrayList<>();
PTransformRunnerFactoryTestContext context =
PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
+ .beamFnStateClient(fakeClient)
.processBundleInstructionId("57")
.components(
RunnerApi.Components.newBuilder()
@@ -2731,223 +2893,1003 @@ public class FnApiDoFnRunnerTest implements Serializable {
.putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
.build())
.build();
+ List<WindowedValue<KV<KV<String, OffsetRange>, Double>>> mainOutputValues = new ArrayList<>();
+ Coder coder =
+ KvCoder.of(
+ KvCoder.of(
+ StringUtf8Coder.of(), KvCoder.of(OffsetRange.Coder.of(), InstantCoder.of())),
+ DoubleCoder.of());
context.addPCollectionConsumer(
- outputPCollectionId, (FnDataReceiver) new OutputFnDataReceiver(mainOutputValues));
+ outputPCollectionId, (FnDataReceiver) new SplittableFnDataReceiver(mainOutputValues));
new FnApiDoFnRunner.Factory<>().addRunnerForPTransform(context);
+ assertTrue(context.getStartBundleFunctions().isEmpty());
mainOutputValues.clear();
+ assertThat(
+ context.getPCollectionConsumers().keySet(),
+ containsInAnyOrder(inputPCollectionId, outputPCollectionId));
+
FnDataReceiver<WindowedValue<?>> mainInput =
context.getPCollectionConsumer(inputPCollectionId);
- mainInput.accept(valueInGlobalWindow("0"));
- mainInput.accept(timestampedValueInGlobalWindow("3", new Instant(0L)));
+ assertThat(mainInput, instanceOf(HandlesSplits.class));
+
+ mainInput.accept(
+ valueInGlobalWindow(
+ KV.of(
+ KV.of("5", KV.of(new OffsetRange(0, 5), GlobalWindow.TIMESTAMP_MIN_VALUE)),
+ 5.0)));
+ mainInput.accept(
+ valueInGlobalWindow(
+ KV.of(
+ KV.of("2", KV.of(new OffsetRange(0, 2), GlobalWindow.TIMESTAMP_MIN_VALUE)),
+ 2.0)));
+ assertThat(
+ mainOutputValues,
+ contains(
+ valueInGlobalWindow(
+ KV.of(
+ KV.of("5", KV.of(new OffsetRange(0, 2), GlobalWindow.TIMESTAMP_MIN_VALUE)),
+ 2.0)),
+ valueInGlobalWindow(
+ KV.of(
+ KV.of("2", KV.of(new OffsetRange(0, 1), GlobalWindow.TIMESTAMP_MIN_VALUE)),
+ 1.0))));
+ mainOutputValues.clear();
+
+ assertTrue(context.getFinishBundleFunctions().isEmpty());
+ assertThat(mainOutputValues, empty());
+
+ Iterables.getOnlyElement(context.getTearDownFunctions()).run();
+ assertThat(mainOutputValues, empty());
}
- }
- @RunWith(JUnit4.class)
- public static class SplitTest {
- @Rule public final ExpectedException expected = ExpectedException.none();
- private IntervalWindow window1;
- private IntervalWindow window2;
- private IntervalWindow window3;
- private WindowedValue<String> currentElement;
- private OffsetRange currentRestriction;
- private Instant currentWatermarkEstimatorState;
- private Instant initialWatermark;
- KV<Instant, Instant> watermarkAndState;
+ @Test
+ public void testProcessElementForWindowedTruncateAndSizeRestriction() throws Exception {
+ Pipeline p = Pipeline.create();
+ PCollection<String> valuePCollection = p.apply(Create.of("unused"));
+ PCollectionView<String> singletonSideInputView = valuePCollection.apply(View.asSingleton());
+ valuePCollection
+ .apply(Window.into(SlidingWindows.of(Duration.standardSeconds(1))))
+ .apply(
+ TEST_TRANSFORM_ID,
+ ParDo.of(new WindowObservingTestSplittableDoFn(singletonSideInputView))
+ .withSideInputs(singletonSideInputView));
- private static final String PROCESS_TRANSFORM_ID = "processPTransformId";
- private static final String TRUNCATE_TRANSFORM_ID = "truncatePTransformId";
- private static final String PROCESS_INPUT_ID = "processInputId";
- private static final String TRUNCATE_INPUT_ID = "truncateInputId";
- private static final String PROCESS_OUTPUT_ID = "processOutputId";
+ RunnerApi.Pipeline pProto =
+ ProtoOverrides.updateTransform(
+ PTransformTranslation.PAR_DO_TRANSFORM_URN,
+ PipelineTranslation.toProto(p, SdkComponents.create(p.getOptions()), true),
+ SplittableParDoExpander.createTruncateReplacement());
+ String expandedTransformId =
+ Iterables.find(
+ pProto.getComponents().getTransformsMap().entrySet(),
+ entry ->
+ entry
+ .getValue()
+ .getSpec()
+ .getUrn()
+ .equals(
+ PTransformTranslation.SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN)
+ && entry.getValue().getUniqueName().contains(TEST_TRANSFORM_ID))
+ .getKey();
+ RunnerApi.PTransform pTransform =
+ pProto.getComponents().getTransformsOrThrow(expandedTransformId);
+ String inputPCollectionId =
+ pTransform.getInputsOrThrow(ParDoTranslation.getMainInputName(pTransform));
+ String outputPCollectionId = Iterables.getOnlyElement(pTransform.getOutputsMap().values());
- private KV<WindowedValue, WindowedValue> createSplitInWindow(
- OffsetRange primaryRestriction, OffsetRange residualRestriction, BoundedWindow window) {
- return KV.of(
- WindowedValues.of(
+ FakeBeamFnStateClient fakeClient =
+ new FakeBeamFnStateClient(StringUtf8Coder.of(), ImmutableMap.of());
+
+ PTransformRunnerFactoryTestContext context =
+ PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
+ .beamFnStateClient(fakeClient)
+ .processBundleInstructionId("57")
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .putAllEnvironments(Collections.emptyMap())
+ .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .build())
+ .build();
+ List<WindowedValue<KV<KV<String, OffsetRange>, Double>>> mainOutputValues = new ArrayList<>();
+ Coder coder =
+ KvCoder.of(
+ KvCoder.of(
+ StringUtf8Coder.of(), KvCoder.of(OffsetRange.Coder.of(), InstantCoder.of())),
+ DoubleCoder.of());
+ context.addPCollectionConsumer(
+ outputPCollectionId, (FnDataReceiver) new SplittableFnDataReceiver(mainOutputValues));
+
+ new FnApiDoFnRunner.Factory<>().addRunnerForPTransform(context);
+
+ assertTrue(context.getStartBundleFunctions().isEmpty());
+ mainOutputValues.clear();
+
+ assertThat(
+ context.getPCollectionConsumers().keySet(),
+ containsInAnyOrder(inputPCollectionId, outputPCollectionId));
+
+ FnDataReceiver<WindowedValue<?>> mainInput =
+ context.getPCollectionConsumer(inputPCollectionId);
+ assertThat(mainInput, instanceOf(HandlesSplits.class));
+
+ IntervalWindow window1 = new IntervalWindow(new Instant(5), new Instant(10));
+ IntervalWindow window2 = new IntervalWindow(new Instant(6), new Instant(11));
+ WindowedValue<?> firstValue =
+ valueInWindows(
KV.of(
- currentElement.getValue(),
- KV.of(primaryRestriction, currentWatermarkEstimatorState)),
- currentElement.getTimestamp(),
- window,
- currentElement.getPaneInfo()),
- WindowedValues.of(
+ KV.of("5", KV.of(new OffsetRange(0, 5), GlobalWindow.TIMESTAMP_MIN_VALUE)), 5.0),
+ window1,
+ window2);
+ WindowedValue<?> secondValue =
+ valueInWindows(
KV.of(
- currentElement.getValue(),
- KV.of(residualRestriction, watermarkAndState.getValue())),
- currentElement.getTimestamp(),
- window,
- currentElement.getPaneInfo()));
- }
-
- private KV<WindowedValue, WindowedValue> createSplitAcrossWindows(
- List<BoundedWindow> primaryWindows, List<BoundedWindow> residualWindows) {
- return KV.of(
- primaryWindows.isEmpty()
- ? null
- : WindowedValues.of(
+ KV.of("2", KV.of(new OffsetRange(0, 2), GlobalWindow.TIMESTAMP_MIN_VALUE)), 2.0),
+ window1,
+ window2);
+ mainInput.accept(firstValue);
+ mainInput.accept(secondValue);
+ assertThat(
+ mainOutputValues,
+ contains(
+ WindowedValues.of(
KV.of(
- currentElement.getValue(),
- KV.of(currentRestriction, currentWatermarkEstimatorState)),
- currentElement.getTimestamp(),
- primaryWindows,
- currentElement.getPaneInfo()),
- residualWindows.isEmpty()
- ? null
- : WindowedValues.of(
+ KV.of("5", KV.of(new OffsetRange(0, 2), GlobalWindow.TIMESTAMP_MIN_VALUE)),
+ 2.0),
+ firstValue.getTimestamp(),
+ window1,
+ firstValue.getPaneInfo()),
+ WindowedValues.of(
KV.of(
- currentElement.getValue(),
- KV.of(currentRestriction, currentWatermarkEstimatorState)),
- currentElement.getTimestamp(),
- residualWindows,
- currentElement.getPaneInfo()));
- }
-
- private KV<WindowedValue, WindowedValue> createSplitWithSizeInWindow(
- OffsetRange primaryRestriction, OffsetRange residualRestriction, BoundedWindow window) {
- return KV.of(
- WindowedValues.of(
- KV.of(
+ KV.of("5", KV.of(new OffsetRange(0, 2), GlobalWindow.TIMESTAMP_MIN_VALUE)),
+ 2.0),
+ firstValue.getTimestamp(),
+ window2,
+ firstValue.getPaneInfo()),
+ WindowedValues.of(
KV.of(
- currentElement.getValue(),
- KV.of(primaryRestriction, currentWatermarkEstimatorState)),
- (double) (primaryRestriction.getTo() - primaryRestriction.getFrom())),
- currentElement.getTimestamp(),
- window,
- currentElement.getPaneInfo()),
- WindowedValues.of(
- KV.of(
+ KV.of("2", KV.of(new OffsetRange(0, 1), GlobalWindow.TIMESTAMP_MIN_VALUE)),
+ 1.0),
+ firstValue.getTimestamp(),
+ window1,
+ firstValue.getPaneInfo()),
+ WindowedValues.of(
KV.of(
- currentElement.getValue(),
- KV.of(residualRestriction, watermarkAndState.getValue())),
- (double) (residualRestriction.getTo() - residualRestriction.getFrom())),
- currentElement.getTimestamp(),
- window,
- currentElement.getPaneInfo()));
+ KV.of("2", KV.of(new OffsetRange(0, 1), GlobalWindow.TIMESTAMP_MIN_VALUE)),
+ 1.0),
+ firstValue.getTimestamp(),
+ window2,
+ firstValue.getPaneInfo())));
+ mainOutputValues.clear();
+
+ assertTrue(context.getFinishBundleFunctions().isEmpty());
+ assertThat(mainOutputValues, empty());
+
+ Iterables.getOnlyElement(context.getTearDownFunctions()).run();
+ assertThat(mainOutputValues, empty());
}
- private KV<WindowedValue, WindowedValue> createSplitWithSizeAcrossWindows(
- List<BoundedWindow> primaryWindows, List<BoundedWindow> residualWindows) {
- return KV.of(
- primaryWindows.isEmpty()
- ? null
- : WindowedValues.of(
+ @Test
+ public void
+ testProcessElementForWindowedTruncateAndSizeRestrictionWithNonWindowObservingOptimization()
+ throws Exception {
+ Pipeline p = Pipeline.create();
+ PCollection<String> valuePCollection = p.apply(Create.of("unused"));
+ valuePCollection
+ .apply(Window.into(SlidingWindows.of(Duration.standardSeconds(1))))
+ .apply(TEST_TRANSFORM_ID, ParDo.of(new NonWindowObservingTestSplittableDoFn()));
+
+ RunnerApi.Pipeline pProto =
+ ProtoOverrides.updateTransform(
+ PTransformTranslation.PAR_DO_TRANSFORM_URN,
+ PipelineTranslation.toProto(p, SdkComponents.create(p.getOptions()), true),
+ SplittableParDoExpander.createTruncateReplacement());
+ String expandedTransformId =
+ Iterables.find(
+ pProto.getComponents().getTransformsMap().entrySet(),
+ entry ->
+ entry
+ .getValue()
+ .getSpec()
+ .getUrn()
+ .equals(
+ PTransformTranslation.SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN)
+ && entry.getValue().getUniqueName().contains(TEST_TRANSFORM_ID))
+ .getKey();
+ RunnerApi.PTransform pTransform =
+ pProto.getComponents().getTransformsOrThrow(expandedTransformId);
+ String inputPCollectionId =
+ pTransform.getInputsOrThrow(ParDoTranslation.getMainInputName(pTransform));
+ String outputPCollectionId = Iterables.getOnlyElement(pTransform.getOutputsMap().values());
+
+ PTransformRunnerFactoryTestContext context =
+ PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
+ .processBundleInstructionId("57")
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .putAllEnvironments(Collections.emptyMap())
+ .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .build())
+ .build();
+ List<WindowedValue<KV<KV<String, OffsetRange>, Double>>> mainOutputValues = new ArrayList<>();
+ Coder coder =
+ KvCoder.of(
+ KvCoder.of(
+ StringUtf8Coder.of(), KvCoder.of(OffsetRange.Coder.of(), InstantCoder.of())),
+ DoubleCoder.of());
+ context.addPCollectionConsumer(
+ outputPCollectionId, (FnDataReceiver) new SplittableFnDataReceiver(mainOutputValues));
+
+ new FnApiDoFnRunner.Factory<>().addRunnerForPTransform(context);
+
+ assertTrue(context.getStartBundleFunctions().isEmpty());
+ mainOutputValues.clear();
+
+ assertThat(
+ context.getPCollectionConsumers().keySet(),
+ containsInAnyOrder(inputPCollectionId, outputPCollectionId));
+
+ FnDataReceiver<WindowedValue<?>> mainInput =
+ context.getPCollectionConsumer(inputPCollectionId);
+ assertThat(mainInput, instanceOf(HandlesSplits.class));
+
+ IntervalWindow window1 = new IntervalWindow(new Instant(5), new Instant(10));
+ IntervalWindow window2 = new IntervalWindow(new Instant(6), new Instant(11));
+ WindowedValue<?> firstValue =
+ valueInWindows(
+ KV.of(
+ KV.of("5", KV.of(new OffsetRange(0, 5), GlobalWindow.TIMESTAMP_MIN_VALUE)), 5.0),
+ window1,
+ window2);
+ WindowedValue<?> secondValue =
+ valueInWindows(
+ KV.of(
+ KV.of("2", KV.of(new OffsetRange(0, 2), GlobalWindow.TIMESTAMP_MIN_VALUE)), 2.0),
+ window1,
+ window2);
+ mainInput.accept(firstValue);
+ mainInput.accept(secondValue);
+ // Ensure that each output element is in all the windows and not one per window.
+ assertThat(
+ mainOutputValues,
+ contains(
+ WindowedValues.of(
KV.of(
- KV.of(
- currentElement.getValue(),
- KV.of(currentRestriction, currentWatermarkEstimatorState)),
- (double) (currentRestriction.getTo() - currentRestriction.getFrom())),
- currentElement.getTimestamp(),
- primaryWindows,
- currentElement.getPaneInfo()),
- residualWindows.isEmpty()
- ? null
- : WindowedValues.of(
+ KV.of("5", KV.of(new OffsetRange(0, 2), GlobalWindow.TIMESTAMP_MIN_VALUE)),
+ 2.0),
+ firstValue.getTimestamp(),
+ ImmutableList.of(window1, window2),
+ firstValue.getPaneInfo()),
+ WindowedValues.of(
KV.of(
- KV.of(
- currentElement.getValue(),
- KV.of(currentRestriction, currentWatermarkEstimatorState)),
- (double) (currentRestriction.getTo() - currentRestriction.getFrom())),
- currentElement.getTimestamp(),
- residualWindows,
- currentElement.getPaneInfo()));
- }
+ KV.of("2", KV.of(new OffsetRange(0, 1), GlobalWindow.TIMESTAMP_MIN_VALUE)),
+ 1.0),
+ firstValue.getTimestamp(),
+ ImmutableList.of(window1, window2),
+ firstValue.getPaneInfo())));
+ mainOutputValues.clear();
- @Before
- public void setUp() {
- window1 = new IntervalWindow(Instant.ofEpochMilli(0), Instant.ofEpochMilli(10));
- window2 = new IntervalWindow(Instant.ofEpochMilli(10), Instant.ofEpochMilli(20));
- window3 = new IntervalWindow(Instant.ofEpochMilli(20), Instant.ofEpochMilli(30));
- currentElement =
- WindowedValues.of(
- "a",
- Instant.ofEpochMilli(57),
- ImmutableList.of(window1, window2, window3),
- PaneInfo.NO_FIRING);
- currentRestriction = new OffsetRange(0L, 100L);
- currentWatermarkEstimatorState = Instant.ofEpochMilli(21);
- initialWatermark = Instant.ofEpochMilli(25);
- watermarkAndState = KV.of(Instant.ofEpochMilli(42), Instant.ofEpochMilli(42));
+ assertTrue(context.getFinishBundleFunctions().isEmpty());
+ assertThat(mainOutputValues, empty());
+
+ Iterables.getOnlyElement(context.getTearDownFunctions()).run();
+ assertThat(mainOutputValues, empty());
}
- @Test
- public void testComputeSplitForProcessWithNullTrackerAndSplitDelegate() throws Exception {
- expected.expect(IllegalArgumentException.class);
- FnApiDoFnRunner.computeSplitForProcess(
- currentElement,
- currentRestriction,
- window1,
- ImmutableList.copyOf(currentElement.getWindows()),
- currentWatermarkEstimatorState,
- 0.0,
- null,
- null,
- null,
+ /**
+ * A {@link DoFn} that outputs elements with timestamp equal to the input timestamp minus the
+ * input element.
+ */
+ private static class SkewingDoFn extends DoFn<String, String> {
+ private final Duration allowedSkew;
+
+ private SkewingDoFn(Duration allowedSkew) {
+ this.allowedSkew = allowedSkew;
+ }
+
+ @ProcessElement
+ public void processElement(ProcessContext context) {
+ Duration duration = Duration.millis(Long.valueOf(context.element()));
+ context.outputWithTimestamp(context.element(), context.timestamp().minus(duration));
+ }
+
+ @Override
+ public Duration getAllowedTimestampSkew() {
+ return allowedSkew;
+ }
+ }
+
+ private static class OutputFnDataReceiver implements FnDataReceiver<WindowedValue> {
+ OutputFnDataReceiver(List<WindowedValue<String>> mainOutputValues) {
+ this.mainOutputValues = mainOutputValues;
+ }
+
+ private final List<WindowedValue<String>> mainOutputValues;
+
+ @Override
+ public void accept(WindowedValue input) throws Exception {
+ mainOutputValues.add(input);
+ }
+ }
+
+ @Test
+ public void testDoFnSkewNotAllowed() throws Exception {
+ Pipeline p = Pipeline.create();
+ PCollection<String> valuePCollection = p.apply(Create.of("0", "1"));
+ PCollection<String> outputPCollection =
+ valuePCollection.apply(TEST_TRANSFORM_ID, ParDo.of(new SkewingDoFn(Duration.ZERO)));
+
+ SdkComponents sdkComponents = SdkComponents.create(p.getOptions());
+ RunnerApi.Pipeline pProto = PipelineTranslation.toProto(p, sdkComponents);
+ String inputPCollectionId = sdkComponents.registerPCollection(valuePCollection);
+ String outputPCollectionId = sdkComponents.registerPCollection(outputPCollection);
+ RunnerApi.PTransform pTransform =
+ pProto
+ .getComponents()
+ .getTransformsOrThrow(
+ pProto
+ .getComponents()
+ .getTransformsOrThrow(TEST_TRANSFORM_ID)
+ .getSubtransforms(0));
+
+ PTransformRunnerFactoryTestContext context =
+ PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
+ .processBundleInstructionId("57")
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .putAllEnvironments(Collections.emptyMap())
+ .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .build())
+ .build();
+ List<WindowedValue<String>> mainOutputValues = new ArrayList<>();
+ Coder coder = StringUtf8Coder.of();
+ context.addPCollectionConsumer(
+ outputPCollectionId, (FnDataReceiver) new OutputFnDataReceiver(mainOutputValues));
+
+ new FnApiDoFnRunner.Factory<>().addRunnerForPTransform(context);
+
+ mainOutputValues.clear();
+ FnDataReceiver<WindowedValue<?>> mainInput =
+ context.getPCollectionConsumer(inputPCollectionId);
+ mainInput.accept(valueInGlobalWindow("0"));
+
+ String message =
+ assertThrows(
+ UserCodeException.class,
+ () -> {
+ mainInput.accept(timestampedValueInGlobalWindow("1", new Instant(0L)));
+ })
+ .getMessage();
+
+ assertThat(
+ message,
+ allOf(
+ containsString(
+ String.format("timestamp %s", new Instant(0).minus(Duration.millis(1L)))),
+ containsString(
+ String.format(
+ "allowed skew (%s)",
+ PeriodFormat.getDefault().print(Duration.ZERO.toPeriod())))));
+ }
+
+ @Test
+ public void testDoFnSkewAllowed() throws Exception {
+ Pipeline p = Pipeline.create();
+ PCollection<String> valuePCollection = p.apply(Create.of("0", "3"));
+ PCollection<String> outputPCollection =
+ valuePCollection.apply(TEST_TRANSFORM_ID, ParDo.of(new SkewingDoFn(Duration.millis(5L))));
+
+ SdkComponents sdkComponents = SdkComponents.create(p.getOptions());
+ RunnerApi.Pipeline pProto = PipelineTranslation.toProto(p, sdkComponents);
+ String inputPCollectionId = sdkComponents.registerPCollection(valuePCollection);
+ String outputPCollectionId = sdkComponents.registerPCollection(outputPCollection);
+ RunnerApi.PTransform pTransform =
+ pProto
+ .getComponents()
+ .getTransformsOrThrow(
+ pProto
+ .getComponents()
+ .getTransformsOrThrow(TEST_TRANSFORM_ID)
+ .getSubtransforms(0));
+
+ List<WindowedValue<String>> mainOutputValues = new ArrayList<>();
+ PTransformRunnerFactoryTestContext context =
+ PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
+ .processBundleInstructionId("57")
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .putAllEnvironments(Collections.emptyMap())
+ .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+ .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .build())
+ .build();
+ Coder coder = StringUtf8Coder.of();
+ context.addPCollectionConsumer(
+ outputPCollectionId, (FnDataReceiver) new OutputFnDataReceiver(mainOutputValues));
+
+ new FnApiDoFnRunner.Factory<>().addRunnerForPTransform(context);
+
+ mainOutputValues.clear();
+
+ FnDataReceiver<WindowedValue<?>> mainInput =
+ context.getPCollectionConsumer(inputPCollectionId);
+ mainInput.accept(valueInGlobalWindow("0"));
+ mainInput.accept(timestampedValueInGlobalWindow("3", new Instant(0L)));
+ }
+ }
+
+ @RunWith(JUnit4.class)
+ public static class SplitTest {
+ @Rule public final ExpectedException expected = ExpectedException.none();
+ private IntervalWindow window1;
+ private IntervalWindow window2;
+ private IntervalWindow window3;
+ private WindowedValue<String> currentElement;
+ private OffsetRange currentRestriction;
+ private Instant currentWatermarkEstimatorState;
+ private Instant initialWatermark;
+ KV<Instant, Instant> watermarkAndState;
+
+ private static final String PROCESS_TRANSFORM_ID = "processPTransformId";
+ private static final String TRUNCATE_TRANSFORM_ID = "truncatePTransformId";
+ private static final String PROCESS_INPUT_ID = "processInputId";
+ private static final String TRUNCATE_INPUT_ID = "truncateInputId";
+ private static final String PROCESS_OUTPUT_ID = "processOutputId";
+ private static final String TRUNCATE_OUTPUT_ID = "truncateOutputId";
+
+ private KV<WindowedValue, WindowedValue> createSplitInWindow(
+ OffsetRange primaryRestriction, OffsetRange residualRestriction, BoundedWindow window) {
+ return KV.of(
+ WindowedValues.of(
+ KV.of(
+ currentElement.getValue(),
+ KV.of(primaryRestriction, currentWatermarkEstimatorState)),
+ currentElement.getTimestamp(),
+ window,
+ currentElement.getPaneInfo()),
+ WindowedValues.of(
+ KV.of(
+ currentElement.getValue(),
+ KV.of(residualRestriction, watermarkAndState.getValue())),
+ currentElement.getTimestamp(),
+ window,
+ currentElement.getPaneInfo()));
+ }
+
+ private KV<WindowedValue, WindowedValue> createSplitAcrossWindows(
+ List<BoundedWindow> primaryWindows, List<BoundedWindow> residualWindows) {
+ return KV.of(
+ primaryWindows.isEmpty()
+ ? null
+ : WindowedValues.of(
+ KV.of(
+ currentElement.getValue(),
+ KV.of(currentRestriction, currentWatermarkEstimatorState)),
+ currentElement.getTimestamp(),
+ primaryWindows,
+ currentElement.getPaneInfo()),
+ residualWindows.isEmpty()
+ ? null
+ : WindowedValues.of(
+ KV.of(
+ currentElement.getValue(),
+ KV.of(currentRestriction, currentWatermarkEstimatorState)),
+ currentElement.getTimestamp(),
+ residualWindows,
+ currentElement.getPaneInfo()));
+ }
+
+ private KV<WindowedValue, WindowedValue> createSplitWithSizeInWindow(
+ OffsetRange primaryRestriction, OffsetRange residualRestriction, BoundedWindow window) {
+ return KV.of(
+ WindowedValues.of(
+ KV.of(
+ KV.of(
+ currentElement.getValue(),
+ KV.of(primaryRestriction, currentWatermarkEstimatorState)),
+ (double) (primaryRestriction.getTo() - primaryRestriction.getFrom())),
+ currentElement.getTimestamp(),
+ window,
+ currentElement.getPaneInfo()),
+ WindowedValues.of(
+ KV.of(
+ KV.of(
+ currentElement.getValue(),
+ KV.of(residualRestriction, watermarkAndState.getValue())),
+ (double) (residualRestriction.getTo() - residualRestriction.getFrom())),
+ currentElement.getTimestamp(),
+ window,
+ currentElement.getPaneInfo()));
+ }
+
+ private KV<WindowedValue, WindowedValue> createSplitWithSizeAcrossWindows(
+ List<BoundedWindow> primaryWindows, List<BoundedWindow> residualWindows) {
+ return KV.of(
+ primaryWindows.isEmpty()
+ ? null
+ : WindowedValues.of(
+ KV.of(
+ KV.of(
+ currentElement.getValue(),
+ KV.of(currentRestriction, currentWatermarkEstimatorState)),
+ (double) (currentRestriction.getTo() - currentRestriction.getFrom())),
+ currentElement.getTimestamp(),
+ primaryWindows,
+ currentElement.getPaneInfo()),
+ residualWindows.isEmpty()
+ ? null
+ : WindowedValues.of(
+ KV.of(
+ KV.of(
+ currentElement.getValue(),
+ KV.of(currentRestriction, currentWatermarkEstimatorState)),
+ (double) (currentRestriction.getTo() - currentRestriction.getFrom())),
+ currentElement.getTimestamp(),
+ residualWindows,
+ currentElement.getPaneInfo()));
+ }
+
+ @Before
+ public void setUp() {
+ window1 = new IntervalWindow(Instant.ofEpochMilli(0), Instant.ofEpochMilli(10));
+ window2 = new IntervalWindow(Instant.ofEpochMilli(10), Instant.ofEpochMilli(20));
+ window3 = new IntervalWindow(Instant.ofEpochMilli(20), Instant.ofEpochMilli(30));
+ currentElement =
+ WindowedValues.of(
+ "a",
+ Instant.ofEpochMilli(57),
+ ImmutableList.of(window1, window2, window3),
+ PaneInfo.NO_FIRING);
+ currentRestriction = new OffsetRange(0L, 100L);
+ currentWatermarkEstimatorState = Instant.ofEpochMilli(21);
+ initialWatermark = Instant.ofEpochMilli(25);
+ watermarkAndState = KV.of(Instant.ofEpochMilli(42), Instant.ofEpochMilli(42));
+ }
+
+ @Test
+ public void testScaledProgress() throws Exception {
+ Progress elementProgress = Progress.from(2, 8);
+ // There is only one window.
+ Progress scaledResult = FnApiDoFnRunner.scaleProgress(elementProgress, 0, 1);
+ assertEquals(2, scaledResult.getWorkCompleted(), 0.0);
+ assertEquals(8, scaledResult.getWorkRemaining(), 0.0);
+
+ // We are at the first window of 3 in total.
+ scaledResult = FnApiDoFnRunner.scaleProgress(elementProgress, 0, 3);
+ assertEquals(2, scaledResult.getWorkCompleted(), 0.0);
+ assertEquals(28, scaledResult.getWorkRemaining(), 0.0);
+
+ // We are at the second window of 3 in total.
+ scaledResult = FnApiDoFnRunner.scaleProgress(elementProgress, 1, 3);
+ assertEquals(12, scaledResult.getWorkCompleted(), 0.0);
+ assertEquals(18, scaledResult.getWorkRemaining(), 0.0);
+
+ // We are at the last window of 3 in total.
+ scaledResult = FnApiDoFnRunner.scaleProgress(elementProgress, 2, 3);
+ assertEquals(22, scaledResult.getWorkCompleted(), 0.0);
+ assertEquals(8, scaledResult.getWorkRemaining(), 0.0);
+ }
+
+ @Test
+ public void testComputeSplitForProcessOrTruncateWithNullTrackerAndSplitDelegate()
+ throws Exception {
+ expected.expect(IllegalArgumentException.class);
+ FnApiDoFnRunner.computeSplitForProcessOrTruncate(
+ currentElement,
+ currentRestriction,
+ window1,
+ ImmutableList.copyOf(currentElement.getWindows()),
+ currentWatermarkEstimatorState,
+ 0.0,
+ null,
+ null,
+ null,
+ 0,
+ 3);
+ }
+
+ @Test
+ public void testComputeSplitForProcessOrTruncateWithNotNullTrackerAndDelegate()
+ throws Exception {
+ expected.expect(IllegalArgumentException.class);
+ FnApiDoFnRunner.computeSplitForProcessOrTruncate(
+ currentElement,
+ currentRestriction,
+ window1,
+ ImmutableList.copyOf(currentElement.getWindows()),
+ currentWatermarkEstimatorState,
+ 0.0,
+ new OffsetRangeTracker(currentRestriction),
+ createSplitDelegate(0.3, 0.0, null),
+ null,
0,
3);
}
- @Test
- public void testComputeSplitForProcessWithNotNullTrackerAndDelegate() throws Exception {
- expected.expect(IllegalArgumentException.class);
- FnApiDoFnRunner.computeSplitForProcess(
- currentElement,
- currentRestriction,
- window1,
- ImmutableList.copyOf(currentElement.getWindows()),
- currentWatermarkEstimatorState,
- 0.0,
- new OffsetRangeTracker(currentRestriction),
- createSplitDelegate(0.3, 0.0, null),
- null,
- 0,
- 3);
- }
+ @Test
+ public void testComputeSplitForProcessOrTruncateWithInvalidWatermarkAndState()
+ throws Exception {
+ expected.expect(NullPointerException.class);
+ FnApiDoFnRunner.computeSplitForProcessOrTruncate(
+ currentElement,
+ currentRestriction,
+ window1,
+ ImmutableList.copyOf(currentElement.getWindows()),
+ currentWatermarkEstimatorState,
+ 0.0,
+ new OffsetRangeTracker(currentRestriction),
+ null,
+ null,
+ 0,
+ 3);
+ }
+
+ @Test
+ public void testTrySplitForProcessCheckpointOnFirstWindow() throws Exception {
+ List<BoundedWindow> windows = ImmutableList.copyOf(currentElement.getWindows());
+ OffsetRangeTracker tracker = new OffsetRangeTracker(currentRestriction);
+ tracker.tryClaim(30L);
+ SplitResultsWithStopIndex result =
+ FnApiDoFnRunner.<Instant>computeSplitForProcessOrTruncate(
+ currentElement,
+ currentRestriction,
+ window1,
+ windows,
+ currentWatermarkEstimatorState,
+ 0.0,
+ tracker,
+ null,
+ watermarkAndState,
+ 0,
+ 3);
+ assertEquals(1, result.getNewWindowStopIndex());
+ KV<WindowedValue, WindowedValue> expectedElementSplit =
+ createSplitInWindow(new OffsetRange(0, 31), new OffsetRange(31, 100), window1);
+ KV<WindowedValue, WindowedValue> expectedWindowSplit =
+ createSplitAcrossWindows(ImmutableList.of(), ImmutableList.of(window2, window3));
+ assertEquals(expectedElementSplit.getKey(), result.getWindowSplit().getPrimarySplitRoot());
+ assertEquals(expectedElementSplit.getValue(), result.getWindowSplit().getResidualSplitRoot());
+ assertEquals(
+ expectedWindowSplit.getKey(),
+ result.getWindowSplit().getPrimaryInFullyProcessedWindowsRoot());
+ assertEquals(
+ expectedWindowSplit.getValue(),
+ result.getWindowSplit().getResidualInUnprocessedWindowsRoot());
+ }
+
+ @Test
+ public void testTrySplitForProcessCheckpointOnFirstWindowAfterOneSplit() throws Exception {
+ List<BoundedWindow> windows = ImmutableList.copyOf(currentElement.getWindows());
+ OffsetRangeTracker tracker = new OffsetRangeTracker(currentRestriction);
+ tracker.tryClaim(30L);
+ SplitResultsWithStopIndex result =
+ FnApiDoFnRunner.<Instant>computeSplitForProcessOrTruncate(
+ currentElement,
+ currentRestriction,
+ window1,
+ windows,
+ currentWatermarkEstimatorState,
+ 0.0,
+ tracker,
+ null,
+ watermarkAndState,
+ 0,
+ 2);
+ assertEquals(1, result.getNewWindowStopIndex());
+ KV<WindowedValue, WindowedValue> expectedElementSplit =
+ createSplitInWindow(new OffsetRange(0, 31), new OffsetRange(31, 100), window1);
+ KV<WindowedValue, WindowedValue> expectedWindowSplit =
+ createSplitAcrossWindows(ImmutableList.of(), ImmutableList.of(window2));
+ assertEquals(expectedElementSplit.getKey(), result.getWindowSplit().getPrimarySplitRoot());
+ assertEquals(expectedElementSplit.getValue(), result.getWindowSplit().getResidualSplitRoot());
+ assertEquals(
+ expectedWindowSplit.getKey(),
+ result.getWindowSplit().getPrimaryInFullyProcessedWindowsRoot());
+ assertEquals(
+ expectedWindowSplit.getValue(),
+ result.getWindowSplit().getResidualInUnprocessedWindowsRoot());
+ }
+
+ @Test
+ public void testTrySplitForProcessSplitOnFirstWindow() throws Exception {
+ List<BoundedWindow> windows = ImmutableList.copyOf(currentElement.getWindows());
+ OffsetRangeTracker tracker = new OffsetRangeTracker(currentRestriction);
+ tracker.tryClaim(30L);
+ SplitResultsWithStopIndex result =
+ FnApiDoFnRunner.<Instant>computeSplitForProcessOrTruncate(
+ currentElement,
+ currentRestriction,
+ window1,
+ windows,
+ currentWatermarkEstimatorState,
+ 0.2,
+ tracker,
+ null,
+ watermarkAndState,
+ 0,
+ 3);
+ assertEquals(1, result.getNewWindowStopIndex());
+ KV<WindowedValue, WindowedValue> expectedElementSplit =
+ createSplitInWindow(new OffsetRange(0, 84), new OffsetRange(84, 100), window1);
+ KV<WindowedValue, WindowedValue> expectedWindowSplit =
+ createSplitAcrossWindows(ImmutableList.of(), ImmutableList.of(window2, window3));
+ assertEquals(expectedElementSplit.getKey(), result.getWindowSplit().getPrimarySplitRoot());
+ assertEquals(expectedElementSplit.getValue(), result.getWindowSplit().getResidualSplitRoot());
+ assertEquals(
+ expectedWindowSplit.getKey(),
+ result.getWindowSplit().getPrimaryInFullyProcessedWindowsRoot());
+ assertEquals(
+ expectedWindowSplit.getValue(),
+ result.getWindowSplit().getResidualInUnprocessedWindowsRoot());
+ }
+
+ @Test
+ public void testTrySplitForProcessSplitOnMiddleWindow() throws Exception {
+ List<BoundedWindow> windows = ImmutableList.copyOf(currentElement.getWindows());
+ OffsetRangeTracker tracker = new OffsetRangeTracker(currentRestriction);
+ tracker.tryClaim(30L);
+ SplitResultsWithStopIndex result =
+ FnApiDoFnRunner.computeSplitForProcessOrTruncate(
+ currentElement,
+ currentRestriction,
+ window2,
+ windows,
+ currentWatermarkEstimatorState,
+ 0.2,
+ tracker,
+ null,
+ watermarkAndState,
+ 1,
+ 3);
+ assertEquals(2, result.getNewWindowStopIndex());
+ // Java uses BigDecimal so 0.2 * 170 = 63.9...
+ // BigDecimal.longValue() will round down to 63 instead of the expected 64
+ KV<WindowedValue, WindowedValue> expectedElementSplit =
+ createSplitInWindow(new OffsetRange(0, 63), new OffsetRange(63, 100), window2);
+ KV<WindowedValue, WindowedValue> expectedWindowSplit =
+ createSplitAcrossWindows(ImmutableList.of(window1), ImmutableList.of(window3));
+ assertEquals(expectedElementSplit.getKey(), result.getWindowSplit().getPrimarySplitRoot());
+ assertEquals(expectedElementSplit.getValue(), result.getWindowSplit().getResidualSplitRoot());
+ assertEquals(
+ expectedWindowSplit.getKey(),
+ result.getWindowSplit().getPrimaryInFullyProcessedWindowsRoot());
+ assertEquals(
+ expectedWindowSplit.getValue(),
+ result.getWindowSplit().getResidualInUnprocessedWindowsRoot());
+ }
+
+ @Test
+ public void testTrySplitForProcessSplitOnLastWindow() throws Exception {
+ List<BoundedWindow> windows = ImmutableList.copyOf(currentElement.getWindows());
+ OffsetRangeTracker tracker = new OffsetRangeTracker(currentRestriction);
+ tracker.tryClaim(30L);
+ SplitResultsWithStopIndex result =
+ FnApiDoFnRunner.computeSplitForProcessOrTruncate(
+ currentElement,
+ currentRestriction,
+ window3,
+ windows,
+ currentWatermarkEstimatorState,
+ 0.2,
+ tracker,
+ null,
+ watermarkAndState,
+ 2,
+ 3);
+ assertEquals(3, result.getNewWindowStopIndex());
+ KV<WindowedValue, WindowedValue> expectedElementSplit =
+ createSplitInWindow(new OffsetRange(0, 44), new OffsetRange(44, 100), window3);
+ KV<WindowedValue, WindowedValue> expectedWindowSplit =
+ createSplitAcrossWindows(ImmutableList.of(window1, window2), ImmutableList.of());
+ assertEquals(expectedElementSplit.getKey(), result.getWindowSplit().getPrimarySplitRoot());
+ assertEquals(expectedElementSplit.getValue(), result.getWindowSplit().getResidualSplitRoot());
+ assertEquals(
+ expectedWindowSplit.getKey(),
+ result.getWindowSplit().getPrimaryInFullyProcessedWindowsRoot());
+ assertEquals(
+ expectedWindowSplit.getValue(),
+ result.getWindowSplit().getResidualInUnprocessedWindowsRoot());
+ }
+
+ @Test
+ public void testTrySplitForProcessSplitOnFirstWindowFallback() throws Exception {
+ List<BoundedWindow> windows = ImmutableList.copyOf(currentElement.getWindows());
+ OffsetRangeTracker tracker = new OffsetRangeTracker(currentRestriction);
+ tracker.tryClaim(100L);
+ assertNull(tracker.trySplit(0.0));
+ SplitResultsWithStopIndex result =
+ FnApiDoFnRunner.computeSplitForProcessOrTruncate(
+ currentElement,
+ currentRestriction,
+ window3,
+ windows,
+ currentWatermarkEstimatorState,
+ 0,
+ tracker,
+ null,
+ watermarkAndState,
+ 0,
+ 3);
+ assertEquals(1, result.getNewWindowStopIndex());
+ KV<WindowedValue, WindowedValue> expectedWindowSplit =
+ createSplitAcrossWindows(ImmutableList.of(window1), ImmutableList.of(window2, window3));
+ assertNull(result.getWindowSplit().getPrimarySplitRoot());
+ assertNull(result.getWindowSplit().getResidualSplitRoot());
+ assertEquals(
+ expectedWindowSplit.getKey(),
+ result.getWindowSplit().getPrimaryInFullyProcessedWindowsRoot());
+ assertEquals(
+ expectedWindowSplit.getValue(),
+ result.getWindowSplit().getResidualInUnprocessedWindowsRoot());
+ }
+
+ @Test
+ public void testTrySplitForProcessSplitOnLastWindowWhenNoElementSplit() throws Exception {
+ List<BoundedWindow> windows = ImmutableList.copyOf(currentElement.getWindows());
+ OffsetRangeTracker tracker = new OffsetRangeTracker(currentRestriction);
+ tracker.tryClaim(100L);
+ assertNull(tracker.trySplit(0.0));
+ SplitResultsWithStopIndex result =
+ FnApiDoFnRunner.computeSplitForProcessOrTruncate(
+ currentElement,
+ currentRestriction,
+ window3,
+ windows,
+ currentWatermarkEstimatorState,
+ 0,
+ tracker,
+ null,
+ watermarkAndState,
+ 2,
+ 3);
+ assertNull(result);
+ }
+
+ @Test
+ public void testTrySplitForProcessOnWindowBoundaryRoundUp() throws Exception {
+ List<BoundedWindow> windows = ImmutableList.copyOf(currentElement.getWindows());
+ OffsetRangeTracker tracker = new OffsetRangeTracker(currentRestriction);
+ tracker.tryClaim(30L);
+ SplitResultsWithStopIndex result =
+ FnApiDoFnRunner.computeSplitForProcessOrTruncate(
+ currentElement,
+ currentRestriction,
+ window2,
+ windows,
+ currentWatermarkEstimatorState,
+ 0.6,
+ tracker,
+ null,
+ watermarkAndState,
+ 0,
+ 3);
+ assertEquals(2, result.getNewWindowStopIndex());
+ KV<WindowedValue, WindowedValue> expectedWindowSplit =
+ createSplitAcrossWindows(ImmutableList.of(window1, window2), ImmutableList.of(window3));
+ assertNull(result.getWindowSplit().getPrimarySplitRoot());
+ assertNull(result.getWindowSplit().getResidualSplitRoot());
+ assertEquals(
+ expectedWindowSplit.getKey(),
+ result.getWindowSplit().getPrimaryInFullyProcessedWindowsRoot());
+ assertEquals(
+ expectedWindowSplit.getValue(),
+ result.getWindowSplit().getResidualInUnprocessedWindowsRoot());
+ }
+
+ @Test
+ public void testTrySplitForProcessOnWindowBoundaryRoundDown() throws Exception {
+ List<BoundedWindow> windows = ImmutableList.copyOf(currentElement.getWindows());
+ OffsetRangeTracker tracker = new OffsetRangeTracker(currentRestriction);
+ tracker.tryClaim(30L);
+ SplitResultsWithStopIndex result =
+ FnApiDoFnRunner.computeSplitForProcessOrTruncate(
+ currentElement,
+ currentRestriction,
+ window2,
+ windows,
+ currentWatermarkEstimatorState,
+ 0.3,
+ tracker,
+ null,
+ watermarkAndState,
+ 0,
+ 3);
+ assertEquals(1, result.getNewWindowStopIndex());
+ KV<WindowedValue, WindowedValue> expectedWindowSplit =
+ createSplitAcrossWindows(ImmutableList.of(window1), ImmutableList.of(window2, window3));
+ assertNull(result.getWindowSplit().getPrimarySplitRoot());
+ assertNull(result.getWindowSplit().getResidualSplitRoot());
+ assertEquals(
+ expectedWindowSplit.getKey(),
+ result.getWindowSplit().getPrimaryInFullyProcessedWindowsRoot());
+ assertEquals(
+ expectedWindowSplit.getValue(),
+ result.getWindowSplit().getResidualInUnprocessedWindowsRoot());
+ }
+
+ @Test
+ public void testTrySplitForProcessOnWindowBoundaryRoundDownOnLastWindow() throws Exception {
+ List<BoundedWindow> windows = ImmutableList.copyOf(currentElement.getWindows());
+ OffsetRangeTracker tracker = new OffsetRangeTracker(currentRestriction);
+ tracker.tryClaim(30L);
+ SplitResultsWithStopIndex result =
+ FnApiDoFnRunner.computeSplitForProcessOrTruncate(
+ currentElement,
+ currentRestriction,
+ window2,
+ windows,
+ currentWatermarkEstimatorState,
+ 0.9,
+ tracker,
+ null,
+ watermarkAndState,
+ 0,
+ 3);
+ assertEquals(2, result.getNewWindowStopIndex());
+ KV<WindowedValue, WindowedValue> expectedWindowSplit =
+ createSplitAcrossWindows(ImmutableList.of(window1, window2), ImmutableList.of(window3));
+ assertNull(result.getWindowSplit().getPrimarySplitRoot());
+ assertNull(result.getWindowSplit().getResidualSplitRoot());
+ assertEquals(
+ expectedWindowSplit.getKey(),
+ result.getWindowSplit().getPrimaryInFullyProcessedWindowsRoot());
+ assertEquals(
+ expectedWindowSplit.getValue(),
+ result.getWindowSplit().getResidualInUnprocessedWindowsRoot());
+ }
+
+ private HandlesSplits createSplitDelegate(
+ double progress, double expectedFraction, HandlesSplits.SplitResult result) {
+ return new HandlesSplits() {
+ @Override
+ public SplitResult trySplit(double fractionOfRemainder) {
+ checkArgument(fractionOfRemainder == expectedFraction);
+ return result;
+ }
- @Test
- public void testComputeSplitForProcessWithInvalidWatermarkAndState() throws Exception {
- expected.expect(NullPointerException.class);
- FnApiDoFnRunner.computeSplitForProcess(
- currentElement,
- currentRestriction,
- window1,
- ImmutableList.copyOf(currentElement.getWindows()),
- currentWatermarkEstimatorState,
- 0.0,
- new OffsetRangeTracker(currentRestriction),
- null,
- null,
- 0,
- 3);
+ @Override
+ public double getProgress() {
+ return progress;
+ }
+ };
}
@Test
- public void testTrySplitForProcessCheckpointOnFirstWindow() throws Exception {
+ public void testTrySplitForTruncateCheckpointOnFirstWindow() throws Exception {
List<BoundedWindow> windows = ImmutableList.copyOf(currentElement.getWindows());
- OffsetRangeTracker tracker = new OffsetRangeTracker(currentRestriction);
- tracker.tryClaim(30L);
+ SplitResult splitResult =
+ SplitResult.of(
+ ImmutableList.of(BundleApplication.getDefaultInstance()),
+ ImmutableList.of(DelayedBundleApplication.getDefaultInstance()));
+ HandlesSplits splitDelegate = createSplitDelegate(0.3, 0.0, splitResult);
SplitResultsWithStopIndex result =
- FnApiDoFnRunner.<Instant>computeSplitForProcess(
+ FnApiDoFnRunner.computeSplitForProcessOrTruncate(
currentElement,
currentRestriction,
window1,
windows,
currentWatermarkEstimatorState,
0.0,
- tracker,
null,
- watermarkAndState,
+ splitDelegate,
+ null,
0,
3);
assertEquals(1, result.getNewWindowStopIndex());
- KV<WindowedValue, WindowedValue> expectedElementSplit =
- createSplitInWindow(new OffsetRange(0, 31), new OffsetRange(31, 100), window1);
KV<WindowedValue, WindowedValue> expectedWindowSplit =
createSplitAcrossWindows(ImmutableList.of(), ImmutableList.of(window2, window3));
- assertEquals(expectedElementSplit.getKey(), result.getWindowSplit().getPrimarySplitRoot());
- assertEquals(expectedElementSplit.getValue(), result.getWindowSplit().getResidualSplitRoot());
+ assertEquals(splitResult, result.getDownstreamSplit());
+ assertNull(result.getWindowSplit().getPrimarySplitRoot());
+ assertNull(result.getWindowSplit().getResidualSplitRoot());
assertEquals(
expectedWindowSplit.getKey(),
result.getWindowSplit().getPrimaryInFullyProcessedWindowsRoot());
@@ -2957,30 +3899,32 @@ public class FnApiDoFnRunnerTest implements Serializable {
}
@Test
- public void testTrySplitForProcessCheckpointOnFirstWindowAfterOneSplit() throws Exception {
+ public void testTrySplitForTruncateCheckpointOnFirstWindowAfterOneSplit() throws Exception {
List<BoundedWindow> windows = ImmutableList.copyOf(currentElement.getWindows());
- OffsetRangeTracker tracker = new OffsetRangeTracker(currentRestriction);
- tracker.tryClaim(30L);
+ SplitResult splitResult =
+ SplitResult.of(
+ ImmutableList.of(BundleApplication.getDefaultInstance()),
+ ImmutableList.of(DelayedBundleApplication.getDefaultInstance()));
+ HandlesSplits splitDelegate = createSplitDelegate(0.3, 0.0, splitResult);
SplitResultsWithStopIndex result =
- FnApiDoFnRunner.<Instant>computeSplitForProcess(
+ FnApiDoFnRunner.computeSplitForProcessOrTruncate(
currentElement,
currentRestriction,
window1,
windows,
currentWatermarkEstimatorState,
0.0,
- tracker,
null,
- watermarkAndState,
+ splitDelegate,
+ null,
0,
2);
assertEquals(1, result.getNewWindowStopIndex());
- KV<WindowedValue, WindowedValue> expectedElementSplit =
- createSplitInWindow(new OffsetRange(0, 31), new OffsetRange(31, 100), window1);
KV<WindowedValue, WindowedValue> expectedWindowSplit =
createSplitAcrossWindows(ImmutableList.of(), ImmutableList.of(window2));
- assertEquals(expectedElementSplit.getKey(), result.getWindowSplit().getPrimarySplitRoot());
- assertEquals(expectedElementSplit.getValue(), result.getWindowSplit().getResidualSplitRoot());
+ assertEquals(splitResult, result.getDownstreamSplit());
+ assertNull(result.getWindowSplit().getPrimarySplitRoot());
+ assertNull(result.getWindowSplit().getResidualSplitRoot());
assertEquals(
expectedWindowSplit.getKey(),
result.getWindowSplit().getPrimaryInFullyProcessedWindowsRoot());
@@ -2990,30 +3934,32 @@ public class FnApiDoFnRunnerTest implements Serializable {
}
@Test
- public void testTrySplitForProcessSplitOnFirstWindow() throws Exception {
+ public void testTrySplitForTruncateSplitOnFirstWindow() throws Exception {
List<BoundedWindow> windows = ImmutableList.copyOf(currentElement.getWindows());
- OffsetRangeTracker tracker = new OffsetRangeTracker(currentRestriction);
- tracker.tryClaim(30L);
+ SplitResult splitResult =
+ SplitResult.of(
+ ImmutableList.of(BundleApplication.getDefaultInstance()),
+ ImmutableList.of(DelayedBundleApplication.getDefaultInstance()));
+ HandlesSplits splitDelegate = createSplitDelegate(0.3, 0.54, splitResult);
SplitResultsWithStopIndex result =
- FnApiDoFnRunner.<Instant>computeSplitForProcess(
+ FnApiDoFnRunner.computeSplitForProcessOrTruncate(
currentElement,
currentRestriction,
window1,
windows,
currentWatermarkEstimatorState,
0.2,
- tracker,
null,
- watermarkAndState,
+ splitDelegate,
+ null,
0,
3);
assertEquals(1, result.getNewWindowStopIndex());
- KV<WindowedValue, WindowedValue> expectedElementSplit =
- createSplitInWindow(new OffsetRange(0, 84), new OffsetRange(84, 100), window1);
KV<WindowedValue, WindowedValue> expectedWindowSplit =
createSplitAcrossWindows(ImmutableList.of(), ImmutableList.of(window2, window3));
- assertEquals(expectedElementSplit.getKey(), result.getWindowSplit().getPrimarySplitRoot());
- assertEquals(expectedElementSplit.getValue(), result.getWindowSplit().getResidualSplitRoot());
+ assertEquals(splitResult, result.getDownstreamSplit());
+ assertNull(result.getWindowSplit().getPrimarySplitRoot());
+ assertNull(result.getWindowSplit().getResidualSplitRoot());
assertEquals(
expectedWindowSplit.getKey(),
result.getWindowSplit().getPrimaryInFullyProcessedWindowsRoot());
@@ -3023,32 +3969,32 @@ public class FnApiDoFnRunnerTest implements Serializable {
}
@Test
- public void testTrySplitForProcessSplitOnMiddleWindow() throws Exception {
+ public void testTrySplitForTruncateSplitOnMiddleWindow() throws Exception {
List<BoundedWindow> windows = ImmutableList.copyOf(currentElement.getWindows());
- OffsetRangeTracker tracker = new OffsetRangeTracker(currentRestriction);
- tracker.tryClaim(30L);
+ SplitResult splitResult =
+ SplitResult.of(
+ ImmutableList.of(BundleApplication.getDefaultInstance()),
+ ImmutableList.of(DelayedBundleApplication.getDefaultInstance()));
+ HandlesSplits splitDelegate = createSplitDelegate(0.3, 0.34, splitResult);
SplitResultsWithStopIndex result =
- FnApiDoFnRunner.computeSplitForProcess(
+ FnApiDoFnRunner.computeSplitForProcessOrTruncate(
currentElement,
currentRestriction,
- window2,
+ window1,
windows,
currentWatermarkEstimatorState,
0.2,
- tracker,
null,
- watermarkAndState,
+ splitDelegate,
+ null,
1,
3);
assertEquals(2, result.getNewWindowStopIndex());
- // Java uses BigDecimal so 0.2 * 170 = 63.9...
- // BigDecimal.longValue() will round down to 63 instead of the expected 64
- KV<WindowedValue, WindowedValue> expectedElementSplit =
- createSplitInWindow(new OffsetRange(0, 63), new OffsetRange(63, 100), window2);
KV<WindowedValue, WindowedValue> expectedWindowSplit =
createSplitAcrossWindows(ImmutableList.of(window1), ImmutableList.of(window3));
- assertEquals(expectedElementSplit.getKey(), result.getWindowSplit().getPrimarySplitRoot());
- assertEquals(expectedElementSplit.getValue(), result.getWindowSplit().getResidualSplitRoot());
+ assertEquals(splitResult, result.getDownstreamSplit());
+ assertNull(result.getWindowSplit().getPrimarySplitRoot());
+ assertNull(result.getWindowSplit().getResidualSplitRoot());
assertEquals(
expectedWindowSplit.getKey(),
result.getWindowSplit().getPrimaryInFullyProcessedWindowsRoot());
@@ -3058,30 +4004,32 @@ public class FnApiDoFnRunnerTest implements Serializable {
}
@Test
- public void testTrySplitForProcessSplitOnLastWindow() throws Exception {
+ public void testTrySplitForTruncateSplitOnLastWindow() throws Exception {
List<BoundedWindow> windows = ImmutableList.copyOf(currentElement.getWindows());
- OffsetRangeTracker tracker = new OffsetRangeTracker(currentRestriction);
- tracker.tryClaim(30L);
+ SplitResult splitResult =
+ SplitResult.of(
+ ImmutableList.of(BundleApplication.getDefaultInstance()),
+ ImmutableList.of(DelayedBundleApplication.getDefaultInstance()));
+ HandlesSplits splitDelegate = createSplitDelegate(0.3, 0.2, splitResult);
SplitResultsWithStopIndex result =
- FnApiDoFnRunner.computeSplitForProcess(
+ FnApiDoFnRunner.computeSplitForProcessOrTruncate(
currentElement,
currentRestriction,
- window3,
+ window1,
windows,
currentWatermarkEstimatorState,
0.2,
- tracker,
null,
- watermarkAndState,
+ splitDelegate,
+ null,
2,
3);
assertEquals(3, result.getNewWindowStopIndex());
- KV<WindowedValue, WindowedValue> expectedElementSplit =
- createSplitInWindow(new OffsetRange(0, 44), new OffsetRange(44, 100), window3);
KV<WindowedValue, WindowedValue> expectedWindowSplit =
createSplitAcrossWindows(ImmutableList.of(window1, window2), ImmutableList.of());
- assertEquals(expectedElementSplit.getKey(), result.getWindowSplit().getPrimarySplitRoot());
- assertEquals(expectedElementSplit.getValue(), result.getWindowSplit().getResidualSplitRoot());
+ assertEquals(splitResult, result.getDownstreamSplit());
+ assertNull(result.getWindowSplit().getPrimarySplitRoot());
+ assertNull(result.getWindowSplit().getResidualSplitRoot());
assertEquals(
expectedWindowSplit.getKey(),
result.getWindowSplit().getPrimaryInFullyProcessedWindowsRoot());
@@ -3091,27 +4039,30 @@ public class FnApiDoFnRunnerTest implements Serializable {
}
@Test
- public void testTrySplitForProcessSplitOnFirstWindowFallback() throws Exception {
+ public void testTrySplitForTruncateSplitOnFirstWindowFallback() throws Exception {
List<BoundedWindow> windows = ImmutableList.copyOf(currentElement.getWindows());
- OffsetRangeTracker tracker = new OffsetRangeTracker(currentRestriction);
- tracker.tryClaim(100L);
- assertNull(tracker.trySplit(0.0));
+ SplitResult unusedSplitResult =
+ SplitResult.of(
+ ImmutableList.of(BundleApplication.getDefaultInstance()),
+ ImmutableList.of(DelayedBundleApplication.getDefaultInstance()));
+ HandlesSplits splitDelegate = createSplitDelegate(1.0, 0.0, unusedSplitResult);
SplitResultsWithStopIndex result =
- FnApiDoFnRunner.computeSplitForProcess(
+ FnApiDoFnRunner.computeSplitForProcessOrTruncate(
currentElement,
currentRestriction,
- window3,
+ window1,
windows,
currentWatermarkEstimatorState,
- 0,
- tracker,
+ 0.0,
+ null,
+ splitDelegate,
null,
- watermarkAndState,
0,
3);
assertEquals(1, result.getNewWindowStopIndex());
KV<WindowedValue, WindowedValue> expectedWindowSplit =
createSplitAcrossWindows(ImmutableList.of(window1), ImmutableList.of(window2, window3));
+ assertNull(result.getDownstreamSplit());
assertNull(result.getWindowSplit().getPrimarySplitRoot());
assertNull(result.getWindowSplit().getResidualSplitRoot());
assertEquals(
@@ -3123,48 +4074,50 @@ public class FnApiDoFnRunnerTest implements Serializable {
}
@Test
- public void testTrySplitForProcessSplitOnLastWindowWhenNoElementSplit() throws Exception {
+ public void testTrySplitForTruncateSplitOnLastWindowWhenNoElementSplit() throws Exception {
List<BoundedWindow> windows = ImmutableList.copyOf(currentElement.getWindows());
- OffsetRangeTracker tracker = new OffsetRangeTracker(currentRestriction);
- tracker.tryClaim(100L);
- assertNull(tracker.trySplit(0.0));
+ HandlesSplits splitDelegate = createSplitDelegate(1.0, 0.0, null);
SplitResultsWithStopIndex result =
- FnApiDoFnRunner.computeSplitForProcess(
+ FnApiDoFnRunner.computeSplitForProcessOrTruncate(
currentElement,
currentRestriction,
- window3,
+ window1,
windows,
currentWatermarkEstimatorState,
- 0,
- tracker,
+ 0.0,
+ null,
+ splitDelegate,
null,
- watermarkAndState,
2,
3);
assertNull(result);
}
@Test
- public void testTrySplitForProcessOnWindowBoundaryRoundUp() throws Exception {
+ public void testTrySplitForTruncateOnWindowBoundaryRoundUp() throws Exception {
List<BoundedWindow> windows = ImmutableList.copyOf(currentElement.getWindows());
- OffsetRangeTracker tracker = new OffsetRangeTracker(currentRestriction);
- tracker.tryClaim(30L);
+ SplitResult unusedSplitResult =
+ SplitResult.of(
+ ImmutableList.of(BundleApplication.getDefaultInstance()),
+ ImmutableList.of(DelayedBundleApplication.getDefaultInstance()));
+ HandlesSplits splitDelegate = createSplitDelegate(0.3, 0.0, unusedSplitResult);
SplitResultsWithStopIndex result =
- FnApiDoFnRunner.computeSplitForProcess(
+ FnApiDoFnRunner.computeSplitForProcessOrTruncate(
currentElement,
currentRestriction,
- window2,
+ window1,
windows,
currentWatermarkEstimatorState,
0.6,
- tracker,
null,
- watermarkAndState,
+ splitDelegate,
+ null,
0,
3);
assertEquals(2, result.getNewWindowStopIndex());
KV<WindowedValue, WindowedValue> expectedWindowSplit =
createSplitAcrossWindows(ImmutableList.of(window1, window2), ImmutableList.of(window3));
+ assertNull(result.getDownstreamSplit());
assertNull(result.getWindowSplit().getPrimarySplitRoot());
assertNull(result.getWindowSplit().getResidualSplitRoot());
assertEquals(
@@ -3176,26 +4129,30 @@ public class FnApiDoFnRunnerTest implements Serializable {
}
@Test
- public void testTrySplitForProcessOnWindowBoundaryRoundDown() throws Exception {
+ public void testTrySplitForTruncateOnWindowBoundaryRoundDown() throws Exception {
List<BoundedWindow> windows = ImmutableList.copyOf(currentElement.getWindows());
- OffsetRangeTracker tracker = new OffsetRangeTracker(currentRestriction);
- tracker.tryClaim(30L);
+ SplitResult unusedSplitResult =
+ SplitResult.of(
+ ImmutableList.of(BundleApplication.getDefaultInstance()),
+ ImmutableList.of(DelayedBundleApplication.getDefaultInstance()));
+ HandlesSplits splitDelegate = createSplitDelegate(0.3, 0.0, unusedSplitResult);
SplitResultsWithStopIndex result =
- FnApiDoFnRunner.computeSplitForProcess(
+ FnApiDoFnRunner.computeSplitForProcessOrTruncate(
currentElement,
currentRestriction,
- window2,
+ window1,
windows,
currentWatermarkEstimatorState,
0.3,
- tracker,
null,
- watermarkAndState,
+ splitDelegate,
+ null,
0,
3);
assertEquals(1, result.getNewWindowStopIndex());
KV<WindowedValue, WindowedValue> expectedWindowSplit =
createSplitAcrossWindows(ImmutableList.of(window1), ImmutableList.of(window2, window3));
+ assertNull(result.getDownstreamSplit());
assertNull(result.getWindowSplit().getPrimarySplitRoot());
assertNull(result.getWindowSplit().getResidualSplitRoot());
assertEquals(
@@ -3207,26 +4164,30 @@ public class FnApiDoFnRunnerTest implements Serializable {
}
@Test
- public void testTrySplitForProcessOnWindowBoundaryRoundDownOnLastWindow() throws Exception {
+ public void testTrySplitForTruncateOnWindowBoundaryRoundDownOnLastWindow() throws Exception {
List<BoundedWindow> windows = ImmutableList.copyOf(currentElement.getWindows());
- OffsetRangeTracker tracker = new OffsetRangeTracker(currentRestriction);
- tracker.tryClaim(30L);
+ SplitResult unusedSplitResult =
+ SplitResult.of(
+ ImmutableList.of(BundleApplication.getDefaultInstance()),
+ ImmutableList.of(DelayedBundleApplication.getDefaultInstance()));
+ HandlesSplits splitDelegate = createSplitDelegate(0.3, 0.0, unusedSplitResult);
SplitResultsWithStopIndex result =
- FnApiDoFnRunner.computeSplitForProcess(
+ FnApiDoFnRunner.computeSplitForProcessOrTruncate(
currentElement,
currentRestriction,
- window2,
+ window1,
windows,
currentWatermarkEstimatorState,
- 0.9,
- tracker,
+ 0.6,
+ null,
+ splitDelegate,
null,
- watermarkAndState,
0,
3);
assertEquals(2, result.getNewWindowStopIndex());
KV<WindowedValue, WindowedValue> expectedWindowSplit =
createSplitAcrossWindows(ImmutableList.of(window1, window2), ImmutableList.of(window3));
+ assertNull(result.getDownstreamSplit());
assertNull(result.getWindowSplit().getPrimarySplitRoot());
assertNull(result.getWindowSplit().getResidualSplitRoot());
assertEquals(
@@ -3237,30 +4198,14 @@ public class FnApiDoFnRunnerTest implements Serializable {
result.getWindowSplit().getResidualInUnprocessedWindowsRoot());
}
- private HandlesSplits createSplitDelegate(
- double progress, double expectedFraction, HandlesSplits.SplitResult result) {
- return new HandlesSplits() {
- @Override
- public SplitResult trySplit(double fractionOfRemainder) {
- checkArgument(fractionOfRemainder == expectedFraction);
- return result;
- }
-
- @Override
- public double getProgress() {
- return progress;
- }
- };
- }
-
@Test
public void testConstructSplitResultWithInvalidElementSplits() throws Exception {
expected.expect(IllegalArgumentException.class);
FnApiDoFnRunner.constructSplitResult(
WindowedSplitResult.forRoots(
null,
- valueInGlobalWindow("elementPrimary"),
- valueInGlobalWindow("elementResidual"),
+ WindowedValues.valueInGlobalWindow("elementPrimary"),
+ WindowedValues.valueInGlobalWindow("elementResidual"),
null),
HandlesSplits.SplitResult.of(
ImmutableList.of(BundleApplication.getDefaultInstance()),
@@ -3283,6 +4228,24 @@ public class FnApiDoFnRunnerTest implements Serializable {
return WindowedValues.getFullCoder(inputCoder, windowCoder);
}
+ private HandlesSplits.SplitResult getProcessElementSplit(String transformId, String inputId) {
+ return SplitResult.of(
+ ImmutableList.of(
+ BundleApplication.newBuilder()
+ .setTransformId(transformId)
+ .setInputId(inputId)
+ .build()),
+ ImmutableList.of(
+ DelayedBundleApplication.newBuilder()
+ .setApplication(
+ BundleApplication.newBuilder()
+ .setTransformId(transformId)
+ .setInputId(inputId)
+ .build())
+ .setRequestedTimeDelay(Durations.fromMillis(1000L))
+ .build()));
+ }
+
private org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.Timestamp toTimestamp(
Instant time) {
return org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.Timestamp.newBuilder()
@@ -3291,6 +4254,31 @@ public class FnApiDoFnRunnerTest implements Serializable {
.build();
}
+ @Test
+ public void testConstructSplitResultWithElementSplitFromDelegate() throws Exception {
+ Coder fullInputCoder =
+ getFullInputCoder(
+ StringUtf8Coder.of(),
+ OffsetRange.Coder.of(),
+ InstantCoder.of(),
+ IntervalWindow.getCoder());
+ HandlesSplits.SplitResult elementSplit =
+ getProcessElementSplit(PROCESS_TRANSFORM_ID, PROCESS_INPUT_ID);
+ HandlesSplits.SplitResult result =
+ FnApiDoFnRunner.constructSplitResult(
+ null,
+ elementSplit,
+ fullInputCoder,
+ null,
+ null,
+ TRUNCATE_TRANSFORM_ID,
+ TRUNCATE_INPUT_ID,
+ ImmutableList.of(TRUNCATE_OUTPUT_ID),
+ null);
+ assertEquals(elementSplit.getPrimaryRoots(), result.getPrimaryRoots());
+ assertEquals(elementSplit.getResidualRoots(), result.getResidualRoots());
+ }
+
@Test
public void testConstructSplitResultWithElementSplitFromTracker() throws Exception {
Coder fullInputCoder =
@@ -3441,5 +4429,55 @@ public class FnApiDoFnRunnerTest implements Serializable {
elementSplit.getValue(),
fullInputCoder.decode(elementResidual.getApplication().getElement().newInput()));
}
+
+ @Test
+ public void testConstructSplitResultWithElementAndWindowSplitFromTruncate() throws Exception {
+ Coder fullInputCoder =
+ getFullInputCoder(
+ StringUtf8Coder.of(),
+ OffsetRange.Coder.of(),
+ InstantCoder.of(),
+ IntervalWindow.getCoder());
+ KV<WindowedValue, WindowedValue> windowSplit =
+ createSplitWithSizeAcrossWindows(ImmutableList.of(window1), ImmutableList.of(window3));
+ HandlesSplits.SplitResult elementSplit =
+ getProcessElementSplit(PROCESS_TRANSFORM_ID, PROCESS_INPUT_ID);
+ HandlesSplits.SplitResult result =
+ FnApiDoFnRunner.constructSplitResult(
+ WindowedSplitResult.forRoots(
+ windowSplit.getKey(), null, null, windowSplit.getValue()),
+ elementSplit,
+ fullInputCoder,
+ initialWatermark,
+ watermarkAndState,
+ TRUNCATE_TRANSFORM_ID,
+ TRUNCATE_INPUT_ID,
+ ImmutableList.of(TRUNCATE_OUTPUT_ID),
+ Duration.millis(100L));
+ assertEquals(2, result.getPrimaryRoots().size());
+ BundleApplication windowPrimary = result.getPrimaryRoots().get(0);
+ BundleApplication elementPrimary = result.getPrimaryRoots().get(1);
+ assertEquals(TRUNCATE_TRANSFORM_ID, windowPrimary.getTransformId());
+ assertEquals(TRUNCATE_INPUT_ID, windowPrimary.getInputId());
+ assertEquals(
+ windowSplit.getKey(), fullInputCoder.decode(windowPrimary.getElement().newInput()));
+ assertEquals(elementSplit.getPrimaryRoots().get(0), elementPrimary);
+
+ assertEquals(2, result.getResidualRoots().size());
+ DelayedBundleApplication windowResidual = result.getResidualRoots().get(0);
+ DelayedBundleApplication elementResidual = result.getResidualRoots().get(1);
+ assertEquals(
+ org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.Duration.getDefaultInstance(),
+ windowResidual.getRequestedTimeDelay());
+ assertEquals(TRUNCATE_TRANSFORM_ID, windowResidual.getApplication().getTransformId());
+ assertEquals(TRUNCATE_INPUT_ID, windowResidual.getApplication().getInputId());
+ assertEquals(
+ toTimestamp(initialWatermark),
+ windowResidual.getApplication().getOutputWatermarksMap().get(TRUNCATE_OUTPUT_ID));
+ assertEquals(
+ windowSplit.getValue(),
+ fullInputCoder.decode(windowResidual.getApplication().getElement().newInput()));
+ assertEquals(elementSplit.getResidualRoots().get(0), elementResidual);
+ }
}
}
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/ProgressUtilsTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/ProgressUtilsTest.java
deleted file mode 100644
index 313e86a5151..00000000000
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/ProgressUtilsTest.java
+++ /dev/null
@@ -1,54 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.fn.harness;
-
-import static org.junit.Assert.assertEquals;
-
-import java.io.Serializable;
-import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.Progress;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
-
-/** Tests for {@link ProgressUtils}. */
-@RunWith(JUnit4.class)
-public class ProgressUtilsTest implements Serializable {
- @Test
- public void testScaledProgress() throws Exception {
- Progress elementProgress = Progress.from(2, 8);
- // There is only one window.
- Progress scaledResult = ProgressUtils.scaleProgress(elementProgress, 0, 1);
- assertEquals(2, scaledResult.getWorkCompleted(), 0.0);
- assertEquals(8, scaledResult.getWorkRemaining(), 0.0);
-
- // We are at the first window of 3 in total.
- scaledResult = ProgressUtils.scaleProgress(elementProgress, 0, 3);
- assertEquals(2, scaledResult.getWorkCompleted(), 0.0);
- assertEquals(28, scaledResult.getWorkRemaining(), 0.0);
-
- // We are at the second window of 3 in total.
- scaledResult = ProgressUtils.scaleProgress(elementProgress, 1, 3);
- assertEquals(12, scaledResult.getWorkCompleted(), 0.0);
- assertEquals(18, scaledResult.getWorkRemaining(), 0.0);
-
- // We are at the last window of 3 in total.
- scaledResult = ProgressUtils.scaleProgress(elementProgress, 2, 3);
- assertEquals(22, scaledResult.getWorkCompleted(), 0.0);
- assertEquals(8, scaledResult.getWorkRemaining(), 0.0);
- }
-}
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/SplittableTruncateSizedRestrictionsDoFnRunnerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/SplittableTruncateSizedRestrictionsDoFnRunnerTest.java
deleted file mode 100644
index 615a681095e..00000000000
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/SplittableTruncateSizedRestrictionsDoFnRunnerTest.java
+++ /dev/null
@@ -1,1191 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.fn.harness;
-
-import static org.apache.beam.sdk.values.WindowedValues.valueInGlobalWindow;
-import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
-import static org.hamcrest.MatcherAssert.assertThat;
-import static org.hamcrest.Matchers.contains;
-import static org.hamcrest.Matchers.containsInAnyOrder;
-import static org.hamcrest.Matchers.empty;
-import static org.hamcrest.Matchers.instanceOf;
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNull;
-import static org.junit.Assert.assertTrue;
-import static org.junit.Assert.fail;
-
-import java.io.Serializable;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.List;
-import java.util.ServiceLoader;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
-import java.util.concurrent.Future;
-import org.apache.beam.fn.harness.HandlesSplits.SplitResult;
-import org.apache.beam.fn.harness.state.FakeBeamFnStateClient;
-import org.apache.beam.model.fnexecution.v1.BeamFnApi.BundleApplication;
-import org.apache.beam.model.fnexecution.v1.BeamFnApi.DelayedBundleApplication;
-import org.apache.beam.model.pipeline.v1.RunnerApi;
-import org.apache.beam.runners.core.metrics.MetricUpdates.MetricUpdate;
-import org.apache.beam.sdk.Pipeline;
-import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.sdk.coders.DoubleCoder;
-import org.apache.beam.sdk.coders.InstantCoder;
-import org.apache.beam.sdk.coders.KvCoder;
-import org.apache.beam.sdk.coders.StringUtf8Coder;
-import org.apache.beam.sdk.fn.data.FnDataReceiver;
-import org.apache.beam.sdk.io.range.OffsetRange;
-import org.apache.beam.sdk.metrics.MetricKey;
-import org.apache.beam.sdk.metrics.MetricName;
-import org.apache.beam.sdk.testing.ResetDateTimeProvider;
-import org.apache.beam.sdk.transforms.Create;
-import org.apache.beam.sdk.transforms.ParDo;
-import org.apache.beam.sdk.transforms.View;
-import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
-import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
-import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
-import org.apache.beam.sdk.transforms.windowing.PaneInfo;
-import org.apache.beam.sdk.transforms.windowing.SlidingWindows;
-import org.apache.beam.sdk.transforms.windowing.Window;
-import org.apache.beam.sdk.util.ByteStringOutputStream;
-import org.apache.beam.sdk.util.construction.CoderTranslation;
-import org.apache.beam.sdk.util.construction.CoderTranslation.TranslationContext;
-import org.apache.beam.sdk.util.construction.PTransformTranslation;
-import org.apache.beam.sdk.util.construction.ParDoTranslation;
-import org.apache.beam.sdk.util.construction.PipelineTranslation;
-import org.apache.beam.sdk.util.construction.RehydratedComponents;
-import org.apache.beam.sdk.util.construction.SdkComponents;
-import org.apache.beam.sdk.util.construction.graph.ProtoOverrides;
-import org.apache.beam.sdk.util.construction.graph.SplittableParDoExpander;
-import org.apache.beam.sdk.values.KV;
-import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.sdk.values.PCollectionView;
-import org.apache.beam.sdk.values.WindowedValue;
-import org.apache.beam.sdk.values.WindowedValues;
-import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.util.Durations;
-import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
-import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
-import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
-import org.hamcrest.collection.IsMapContaining;
-import org.joda.time.Duration;
-import org.joda.time.Instant;
-import org.junit.Assert;
-import org.junit.Before;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.experimental.runners.Enclosed;
-import org.junit.rules.ExpectedException;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
-
-/** Tests for {@link SplittableTruncateSizedRestrictionsDoFnRunner}. */
-@RunWith(Enclosed.class)
-@SuppressWarnings({
- "rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
- // TODO(https://github.com/apache/beam/issues/21230): Remove when new version of
- // errorprone is released (2.11.0)
- "unused"
-})
-public class SplittableTruncateSizedRestrictionsDoFnRunnerTest implements Serializable {
-
- @RunWith(JUnit4.class)
- public static class ExecutionTest implements Serializable {
- @Rule public transient ResetDateTimeProvider dateTimeProvider = new ResetDateTimeProvider();
-
- public static final String TEST_TRANSFORM_ID = "pTransformId";
-
- /** @return a test MetricUpdate for expected metrics to compare against */
- public MetricUpdate create(String stepName, MetricName name, long value) {
- return MetricUpdate.create(MetricKey.create(stepName, name), value);
- }
-
- private <T> WindowedValue<T> valueInWindows(
- T value, BoundedWindow window, BoundedWindow... windows) {
- return WindowedValues.of(
- value,
- window.maxTimestamp(),
- ImmutableList.<BoundedWindow>builder().add(window).add(windows).build(),
- PaneInfo.NO_FIRING);
- }
-
- @Test
- public void testRegistration() {
- for (PTransformRunnerFactory.Registrar registrar :
- ServiceLoader.load(PTransformRunnerFactory.Registrar.class)) {
- if (registrar instanceof SplittableTruncateSizedRestrictionsDoFnRunner.Registrar) {
- assertThat(
- registrar.getPTransformRunnerFactories(),
- IsMapContaining.hasKey(
- PTransformTranslation.SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN));
- return;
- }
- }
- fail("Expected registrar not found.");
- }
-
- private static SplitResult createSplitResult(double fractionOfRemainder) {
- ByteStringOutputStream primaryBytes = new ByteStringOutputStream();
- ByteStringOutputStream residualBytes = new ByteStringOutputStream();
- try {
- DoubleCoder.of().encode(fractionOfRemainder, primaryBytes);
- DoubleCoder.of().encode(1 - fractionOfRemainder, residualBytes);
- } catch (Exception e) {
- // No-op.
- }
- return SplitResult.of(
- ImmutableList.of(
- BundleApplication.newBuilder()
- .setElement(primaryBytes.toByteString())
- .setInputId("mainInputId-process")
- .setTransformId("processPTransfromId")
- .build()),
- ImmutableList.of(
- DelayedBundleApplication.newBuilder()
- .setApplication(
- BundleApplication.newBuilder()
- .setElement(residualBytes.toByteString())
- .setInputId("mainInputId-process")
- .setTransformId("processPTransfromId")
- .build())
- .build()));
- }
-
- private static class SplittableFnDataReceiver
- implements HandlesSplits, FnDataReceiver<WindowedValue> {
- SplittableFnDataReceiver(
- List<WindowedValue<KV<KV<String, OffsetRange>, Double>>> mainOutputValues) {
- this.mainOutputValues = mainOutputValues;
- }
-
- private final List<WindowedValue<KV<KV<String, OffsetRange>, Double>>> mainOutputValues;
-
- @Override
- public SplitResult trySplit(double fractionOfRemainder) {
- return createSplitResult(fractionOfRemainder);
- }
-
- @Override
- public double getProgress() {
- return 0.7;
- }
-
- @Override
- public void accept(WindowedValue input) throws Exception {
- mainOutputValues.add(input);
- }
- }
-
- @Test
- public void testProcessElementForTruncateAndSizeRestrictionForwardSplitWhenObservingWindows()
- throws Exception {
- Pipeline p = Pipeline.create();
- PCollection<String> valuePCollection = p.apply(Create.of("unused"));
- PCollectionView<String> singletonSideInputView = valuePCollection.apply(View.asSingleton());
- WindowObservingTestSplittableDoFn doFn =
- WindowObservingTestSplittableDoFn.forSplitAtTruncate(singletonSideInputView);
- valuePCollection
- .apply(Window.into(SlidingWindows.of(Duration.standardSeconds(1))))
- .apply(TEST_TRANSFORM_ID, ParDo.of(doFn).withSideInputs(singletonSideInputView));
-
- RunnerApi.Pipeline pProto =
- ProtoOverrides.updateTransform(
- PTransformTranslation.PAR_DO_TRANSFORM_URN,
- PipelineTranslation.toProto(p, SdkComponents.create(p.getOptions()), true),
- SplittableParDoExpander.createTruncateReplacement());
- String expandedTransformId =
- Iterables.find(
- pProto.getComponents().getTransformsMap().entrySet(),
- entry ->
- entry
- .getValue()
- .getSpec()
- .getUrn()
- .equals(
- PTransformTranslation.SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN)
- && entry.getValue().getUniqueName().contains(TEST_TRANSFORM_ID))
- .getKey();
- RunnerApi.PTransform pTransform =
- pProto.getComponents().getTransformsOrThrow(expandedTransformId);
- String inputPCollectionId =
- pTransform.getInputsOrThrow(ParDoTranslation.getMainInputName(pTransform));
- RunnerApi.PCollection inputPCollection =
- pProto.getComponents().getPcollectionsOrThrow(inputPCollectionId);
- RehydratedComponents rehydratedComponents =
- RehydratedComponents.forComponents(pProto.getComponents());
- Coder<WindowedValue> inputCoder =
- WindowedValues.getFullCoder(
- CoderTranslation.fromProto(
- pProto.getComponents().getCodersOrThrow(inputPCollection.getCoderId()),
- rehydratedComponents,
- TranslationContext.DEFAULT),
- (Coder)
- CoderTranslation.fromProto(
- pProto
- .getComponents()
- .getCodersOrThrow(
- pProto
- .getComponents()
- .getWindowingStrategiesOrThrow(
- inputPCollection.getWindowingStrategyId())
- .getWindowCoderId()),
- rehydratedComponents,
- TranslationContext.DEFAULT));
-
- String outputPCollectionId = Iterables.getOnlyElement(pTransform.getOutputsMap().values());
-
- FakeBeamFnStateClient fakeClient =
- new FakeBeamFnStateClient(StringUtf8Coder.of(), ImmutableMap.of());
-
- PTransformRunnerFactoryTestContext context =
- PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
- .beamFnStateClient(fakeClient)
- .processBundleInstructionId("57")
- // .pCollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
- .components(
- RunnerApi.Components.newBuilder()
- .putAllCoders(pProto.getComponents().getCodersMap())
- .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
- .putAllEnvironments(Collections.emptyMap())
- .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
- .build())
- // .coders(pProto.getComponents().getCodersMap())
- // .windowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
- .build();
- List<WindowedValue<KV<KV<String, OffsetRange>, Double>>> mainOutputValues = new ArrayList<>();
- context.addPCollectionConsumer(
- outputPCollectionId, (FnDataReceiver) new SplittableFnDataReceiver(mainOutputValues));
-
- new SplittableTruncateSizedRestrictionsDoFnRunner.Factory().addRunnerForPTransform(context);
- FnDataReceiver<WindowedValue<?>> mainInput =
- context.getPCollectionConsumer(inputPCollectionId);
- assertThat(mainInput, instanceOf(HandlesSplits.class));
-
- mainOutputValues.clear();
- BoundedWindow window1 = new IntervalWindow(new Instant(5), new Instant(10));
- BoundedWindow window2 = new IntervalWindow(new Instant(6), new Instant(11));
- BoundedWindow window3 = new IntervalWindow(new Instant(7), new Instant(12));
- // Setup and launch the trySplit thread.
- ExecutorService executorService = Executors.newSingleThreadExecutor();
- Future<SplitResult> trySplitFuture =
- executorService.submit(
- () -> {
- try {
- doFn.waitForSplitElementToBeProcessed();
- SplitResult result = ((HandlesSplits) mainInput).trySplit(0);
- Assert.assertNotNull(result);
- return result;
- } finally {
- doFn.trySplitPerformed();
- }
- });
-
- WindowedValue<?> splitValue =
- valueInWindows(
- KV.of(
- KV.of("7", KV.of(new OffsetRange(0, 6), GlobalWindow.TIMESTAMP_MIN_VALUE)), 6.0),
- window1,
- window2,
- window3);
- mainInput.accept(splitValue);
- SplitResult trySplitResult = trySplitFuture.get();
-
- // We expect that there are outputs from window1 and window2
- assertThat(
- mainOutputValues,
- contains(
- WindowedValues.of(
- KV.of(
- KV.of("7", KV.of(new OffsetRange(0, 3), GlobalWindow.TIMESTAMP_MIN_VALUE)),
- 3.0),
- splitValue.getTimestamp(),
- window1,
- splitValue.getPaneInfo()),
- WindowedValues.of(
- KV.of(
- KV.of("7", KV.of(new OffsetRange(0, 3), GlobalWindow.TIMESTAMP_MIN_VALUE)),
- 3.0),
- splitValue.getTimestamp(),
- window2,
- splitValue.getPaneInfo())));
-
- SplitResult expectedElementSplit = createSplitResult(0);
- BundleApplication expectedElementSplitPrimary =
- Iterables.getOnlyElement(expectedElementSplit.getPrimaryRoots());
- ByteStringOutputStream primaryBytes = new ByteStringOutputStream();
- inputCoder.encode(
- WindowedValues.of(
- KV.of(
- KV.of("7", KV.of(new OffsetRange(0, 6), GlobalWindow.TIMESTAMP_MIN_VALUE)), 6.0),
- splitValue.getTimestamp(),
- window1,
- splitValue.getPaneInfo()),
- primaryBytes);
- BundleApplication expectedWindowedPrimary =
- BundleApplication.newBuilder()
- .setElement(primaryBytes.toByteString())
- .setInputId(ParDoTranslation.getMainInputName(pTransform))
- .setTransformId(TEST_TRANSFORM_ID)
- .build();
- DelayedBundleApplication expectedElementSplitResidual =
- Iterables.getOnlyElement(expectedElementSplit.getResidualRoots());
- ByteStringOutputStream residualBytes = new ByteStringOutputStream();
- inputCoder.encode(
- WindowedValues.of(
- KV.of(
- KV.of("7", KV.of(new OffsetRange(0, 6), GlobalWindow.TIMESTAMP_MIN_VALUE)), 6.0),
- splitValue.getTimestamp(),
- window3,
- splitValue.getPaneInfo()),
- residualBytes);
- DelayedBundleApplication expectedWindowedResidual =
- DelayedBundleApplication.newBuilder()
- .setApplication(
- BundleApplication.newBuilder()
- .setElement(residualBytes.toByteString())
- .setInputId(ParDoTranslation.getMainInputName(pTransform))
- .setTransformId(TEST_TRANSFORM_ID)
- .build())
- .build();
- assertThat(
- trySplitResult.getPrimaryRoots(),
- containsInAnyOrder(expectedWindowedPrimary, expectedElementSplitPrimary));
- assertThat(
- trySplitResult.getResidualRoots(),
- containsInAnyOrder(expectedWindowedResidual, expectedElementSplitResidual));
- }
-
- @Test
- public void testProcessElementForTruncateAndSizeRestrictionForwardSplitWithoutObservingWindow()
- throws Exception {
- Pipeline p = Pipeline.create();
- PCollection<String> valuePCollection = p.apply(Create.of("unused"));
- valuePCollection.apply(
- TEST_TRANSFORM_ID, ParDo.of(new NonWindowObservingTestSplittableDoFn()));
-
- RunnerApi.Pipeline pProto =
- ProtoOverrides.updateTransform(
- PTransformTranslation.PAR_DO_TRANSFORM_URN,
- PipelineTranslation.toProto(p, SdkComponents.create(p.getOptions()), true),
- SplittableParDoExpander.createTruncateReplacement());
- String expandedTransformId =
- Iterables.find(
- pProto.getComponents().getTransformsMap().entrySet(),
- entry ->
- entry
- .getValue()
- .getSpec()
- .getUrn()
- .equals(
- PTransformTranslation.SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN)
- && entry.getValue().getUniqueName().contains(TEST_TRANSFORM_ID))
- .getKey();
- RunnerApi.PTransform pTransform =
- pProto.getComponents().getTransformsOrThrow(expandedTransformId);
- String inputPCollectionId =
- pTransform.getInputsOrThrow(ParDoTranslation.getMainInputName(pTransform));
- String outputPCollectionId = Iterables.getOnlyElement(pTransform.getOutputsMap().values());
-
- FakeBeamFnStateClient fakeClient =
- new FakeBeamFnStateClient(StringUtf8Coder.of(), ImmutableMap.of());
-
- PTransformRunnerFactoryTestContext context =
- PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
- .beamFnStateClient(fakeClient)
- .processBundleInstructionId("57")
- .components(
- RunnerApi.Components.newBuilder()
- .putAllCoders(pProto.getComponents().getCodersMap())
- .putAllEnvironments(Collections.emptyMap())
- .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
- .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
- .build())
- .build();
- List<WindowedValue<KV<KV<String, OffsetRange>, Double>>> mainOutputValues = new ArrayList<>();
- context.addPCollectionConsumer(
- outputPCollectionId, (FnDataReceiver) new SplittableFnDataReceiver(mainOutputValues));
-
- new SplittableTruncateSizedRestrictionsDoFnRunner.Factory().addRunnerForPTransform(context);
- FnDataReceiver<WindowedValue<?>> mainInput =
- context.getPCollectionConsumer(inputPCollectionId);
- assertThat(mainInput, instanceOf(HandlesSplits.class));
-
- assertEquals(0.7, ((HandlesSplits) mainInput).getProgress(), 0.0);
- assertEquals(createSplitResult(0.4), ((HandlesSplits) mainInput).trySplit(0.4));
- }
-
- @Test
- public void testProcessElementForTruncateAndSizeRestriction() throws Exception {
- Pipeline p = Pipeline.create();
- PCollection<String> valuePCollection = p.apply(Create.of("unused"));
- valuePCollection.apply(
- TEST_TRANSFORM_ID, ParDo.of(new NonWindowObservingTestSplittableDoFn()));
-
- RunnerApi.Pipeline pProto =
- ProtoOverrides.updateTransform(
- PTransformTranslation.PAR_DO_TRANSFORM_URN,
- PipelineTranslation.toProto(p, SdkComponents.create(p.getOptions()), true),
- SplittableParDoExpander.createTruncateReplacement());
- String expandedTransformId =
- Iterables.find(
- pProto.getComponents().getTransformsMap().entrySet(),
- entry ->
- entry
- .getValue()
- .getSpec()
- .getUrn()
- .equals(
- PTransformTranslation.SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN)
- && entry.getValue().getUniqueName().contains(TEST_TRANSFORM_ID))
- .getKey();
- RunnerApi.PTransform pTransform =
- pProto.getComponents().getTransformsOrThrow(expandedTransformId);
- String inputPCollectionId =
- pTransform.getInputsOrThrow(ParDoTranslation.getMainInputName(pTransform));
- String outputPCollectionId = Iterables.getOnlyElement(pTransform.getOutputsMap().values());
-
- FakeBeamFnStateClient fakeClient =
- new FakeBeamFnStateClient(StringUtf8Coder.of(), ImmutableMap.of());
-
- PTransformRunnerFactoryTestContext context =
- PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
- .beamFnStateClient(fakeClient)
- .processBundleInstructionId("57")
- .components(
- RunnerApi.Components.newBuilder()
- .putAllCoders(pProto.getComponents().getCodersMap())
- .putAllEnvironments(Collections.emptyMap())
- .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
- .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
- .build())
- .build();
- List<WindowedValue<KV<KV<String, OffsetRange>, Double>>> mainOutputValues = new ArrayList<>();
- context.addPCollectionConsumer(
- outputPCollectionId, (FnDataReceiver) new SplittableFnDataReceiver(mainOutputValues));
-
- new SplittableTruncateSizedRestrictionsDoFnRunner.Factory().addRunnerForPTransform(context);
-
- assertTrue(context.getStartBundleFunctions().isEmpty());
- mainOutputValues.clear();
-
- assertThat(
- context.getPCollectionConsumers().keySet(),
- containsInAnyOrder(inputPCollectionId, outputPCollectionId));
-
- FnDataReceiver<WindowedValue<?>> mainInput =
- context.getPCollectionConsumer(inputPCollectionId);
- assertThat(mainInput, instanceOf(HandlesSplits.class));
-
- mainInput.accept(
- valueInGlobalWindow(
- KV.of(
- KV.of("5", KV.of(new OffsetRange(0, 5), GlobalWindow.TIMESTAMP_MIN_VALUE)),
- 5.0)));
- mainInput.accept(
- valueInGlobalWindow(
- KV.of(
- KV.of("2", KV.of(new OffsetRange(0, 2), GlobalWindow.TIMESTAMP_MIN_VALUE)),
- 2.0)));
- assertThat(
- mainOutputValues,
- contains(
- valueInGlobalWindow(
- KV.of(
- KV.of("5", KV.of(new OffsetRange(0, 2), GlobalWindow.TIMESTAMP_MIN_VALUE)),
- 2.0)),
- valueInGlobalWindow(
- KV.of(
- KV.of("2", KV.of(new OffsetRange(0, 1), GlobalWindow.TIMESTAMP_MIN_VALUE)),
- 1.0))));
- mainOutputValues.clear();
-
- assertTrue(context.getFinishBundleFunctions().isEmpty());
- assertThat(mainOutputValues, empty());
-
- Iterables.getOnlyElement(context.getTearDownFunctions()).run();
- assertThat(mainOutputValues, empty());
- }
-
- @Test
- public void testProcessElementForWindowedTruncateAndSizeRestriction() throws Exception {
- Pipeline p = Pipeline.create();
- PCollection<String> valuePCollection = p.apply(Create.of("unused"));
- PCollectionView<String> singletonSideInputView = valuePCollection.apply(View.asSingleton());
- valuePCollection
- .apply(Window.into(SlidingWindows.of(Duration.standardSeconds(1))))
- .apply(
- TEST_TRANSFORM_ID,
- ParDo.of(new WindowObservingTestSplittableDoFn(singletonSideInputView))
- .withSideInputs(singletonSideInputView));
-
- RunnerApi.Pipeline pProto =
- ProtoOverrides.updateTransform(
- PTransformTranslation.PAR_DO_TRANSFORM_URN,
- PipelineTranslation.toProto(p, SdkComponents.create(p.getOptions()), true),
- SplittableParDoExpander.createTruncateReplacement());
- String expandedTransformId =
- Iterables.find(
- pProto.getComponents().getTransformsMap().entrySet(),
- entry ->
- entry
- .getValue()
- .getSpec()
- .getUrn()
- .equals(
- PTransformTranslation.SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN)
- && entry.getValue().getUniqueName().contains(TEST_TRANSFORM_ID))
- .getKey();
- RunnerApi.PTransform pTransform =
- pProto.getComponents().getTransformsOrThrow(expandedTransformId);
- String inputPCollectionId =
- pTransform.getInputsOrThrow(ParDoTranslation.getMainInputName(pTransform));
- String outputPCollectionId = Iterables.getOnlyElement(pTransform.getOutputsMap().values());
-
- FakeBeamFnStateClient fakeClient =
- new FakeBeamFnStateClient(StringUtf8Coder.of(), ImmutableMap.of());
-
- PTransformRunnerFactoryTestContext context =
- PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
- .beamFnStateClient(fakeClient)
- .processBundleInstructionId("57")
- .components(
- RunnerApi.Components.newBuilder()
- .putAllCoders(pProto.getComponents().getCodersMap())
- .putAllEnvironments(Collections.emptyMap())
- .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
- .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
- .build())
- .build();
- List<WindowedValue<KV<KV<String, OffsetRange>, Double>>> mainOutputValues = new ArrayList<>();
- context.addPCollectionConsumer(
- outputPCollectionId, (FnDataReceiver) new SplittableFnDataReceiver(mainOutputValues));
-
- new SplittableTruncateSizedRestrictionsDoFnRunner.Factory().addRunnerForPTransform(context);
-
- assertTrue(context.getStartBundleFunctions().isEmpty());
- mainOutputValues.clear();
-
- assertThat(
- context.getPCollectionConsumers().keySet(),
- containsInAnyOrder(inputPCollectionId, outputPCollectionId));
-
- FnDataReceiver<WindowedValue<?>> mainInput =
- context.getPCollectionConsumer(inputPCollectionId);
- assertThat(mainInput, instanceOf(HandlesSplits.class));
-
- IntervalWindow window1 = new IntervalWindow(new Instant(5), new Instant(10));
- IntervalWindow window2 = new IntervalWindow(new Instant(6), new Instant(11));
- WindowedValue<?> firstValue =
- valueInWindows(
- KV.of(
- KV.of("5", KV.of(new OffsetRange(0, 5), GlobalWindow.TIMESTAMP_MIN_VALUE)), 5.0),
- window1,
- window2);
- WindowedValue<?> secondValue =
- valueInWindows(
- KV.of(
- KV.of("2", KV.of(new OffsetRange(0, 2), GlobalWindow.TIMESTAMP_MIN_VALUE)), 2.0),
- window1,
- window2);
- mainInput.accept(firstValue);
- mainInput.accept(secondValue);
- assertThat(
- mainOutputValues,
- contains(
- WindowedValues.of(
- KV.of(
- KV.of("5", KV.of(new OffsetRange(0, 2), GlobalWindow.TIMESTAMP_MIN_VALUE)),
- 2.0),
- firstValue.getTimestamp(),
- window1,
- firstValue.getPaneInfo()),
- WindowedValues.of(
- KV.of(
- KV.of("5", KV.of(new OffsetRange(0, 2), GlobalWindow.TIMESTAMP_MIN_VALUE)),
- 2.0),
- firstValue.getTimestamp(),
- window2,
- firstValue.getPaneInfo()),
- WindowedValues.of(
- KV.of(
- KV.of("2", KV.of(new OffsetRange(0, 1), GlobalWindow.TIMESTAMP_MIN_VALUE)),
- 1.0),
- firstValue.getTimestamp(),
- window1,
- firstValue.getPaneInfo()),
- WindowedValues.of(
- KV.of(
- KV.of("2", KV.of(new OffsetRange(0, 1), GlobalWindow.TIMESTAMP_MIN_VALUE)),
- 1.0),
- firstValue.getTimestamp(),
- window2,
- firstValue.getPaneInfo())));
- mainOutputValues.clear();
-
- assertTrue(context.getFinishBundleFunctions().isEmpty());
- assertThat(mainOutputValues, empty());
-
- Iterables.getOnlyElement(context.getTearDownFunctions()).run();
- assertThat(mainOutputValues, empty());
- }
-
- @Test
- public void
- testProcessElementForWindowedTruncateAndSizeRestrictionWithNonWindowObservingOptimization()
- throws Exception {
- Pipeline p = Pipeline.create();
- PCollection<String> valuePCollection = p.apply(Create.of("unused"));
- valuePCollection
- .apply(Window.into(SlidingWindows.of(Duration.standardSeconds(1))))
- .apply(TEST_TRANSFORM_ID, ParDo.of(new NonWindowObservingTestSplittableDoFn()));
-
- RunnerApi.Pipeline pProto =
- ProtoOverrides.updateTransform(
- PTransformTranslation.PAR_DO_TRANSFORM_URN,
- PipelineTranslation.toProto(p, SdkComponents.create(p.getOptions()), true),
- SplittableParDoExpander.createTruncateReplacement());
- String expandedTransformId =
- Iterables.find(
- pProto.getComponents().getTransformsMap().entrySet(),
- entry ->
- entry
- .getValue()
- .getSpec()
- .getUrn()
- .equals(
- PTransformTranslation.SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN)
- && entry.getValue().getUniqueName().contains(TEST_TRANSFORM_ID))
- .getKey();
- RunnerApi.PTransform pTransform =
- pProto.getComponents().getTransformsOrThrow(expandedTransformId);
- String inputPCollectionId =
- pTransform.getInputsOrThrow(ParDoTranslation.getMainInputName(pTransform));
- String outputPCollectionId = Iterables.getOnlyElement(pTransform.getOutputsMap().values());
-
- PTransformRunnerFactoryTestContext context =
- PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform)
- .processBundleInstructionId("57")
- .components(
- RunnerApi.Components.newBuilder()
- .putAllCoders(pProto.getComponents().getCodersMap())
- .putAllEnvironments(Collections.emptyMap())
- .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
- .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
- .build())
- .build();
- List<WindowedValue<KV<KV<String, OffsetRange>, Double>>> mainOutputValues = new ArrayList<>();
- context.addPCollectionConsumer(
- outputPCollectionId, (FnDataReceiver) new SplittableFnDataReceiver(mainOutputValues));
-
- new SplittableTruncateSizedRestrictionsDoFnRunner.Factory().addRunnerForPTransform(context);
-
- assertTrue(context.getStartBundleFunctions().isEmpty());
- mainOutputValues.clear();
-
- assertThat(
- context.getPCollectionConsumers().keySet(),
- containsInAnyOrder(inputPCollectionId, outputPCollectionId));
-
- FnDataReceiver<WindowedValue<?>> mainInput =
- context.getPCollectionConsumer(inputPCollectionId);
- assertThat(mainInput, instanceOf(HandlesSplits.class));
-
- IntervalWindow window1 = new IntervalWindow(new Instant(5), new Instant(10));
- IntervalWindow window2 = new IntervalWindow(new Instant(6), new Instant(11));
- WindowedValue<?> firstValue =
- valueInWindows(
- KV.of(
- KV.of("5", KV.of(new OffsetRange(0, 5), GlobalWindow.TIMESTAMP_MIN_VALUE)), 5.0),
- window1,
- window2);
- WindowedValue<?> secondValue =
- valueInWindows(
- KV.of(
- KV.of("2", KV.of(new OffsetRange(0, 2), GlobalWindow.TIMESTAMP_MIN_VALUE)), 2.0),
- window1,
- window2);
- mainInput.accept(firstValue);
- mainInput.accept(secondValue);
- // Ensure that each output element is in all the windows and not one per window.
- assertThat(
- mainOutputValues,
- contains(
- WindowedValues.of(
- KV.of(
- KV.of("5", KV.of(new OffsetRange(0, 2), GlobalWindow.TIMESTAMP_MIN_VALUE)),
- 2.0),
- firstValue.getTimestamp(),
- ImmutableList.of(window1, window2),
- firstValue.getPaneInfo()),
- WindowedValues.of(
- KV.of(
- KV.of("2", KV.of(new OffsetRange(0, 1), GlobalWindow.TIMESTAMP_MIN_VALUE)),
- 1.0),
- firstValue.getTimestamp(),
- ImmutableList.of(window1, window2),
- firstValue.getPaneInfo())));
- mainOutputValues.clear();
-
- assertTrue(context.getFinishBundleFunctions().isEmpty());
- assertThat(mainOutputValues, empty());
-
- Iterables.getOnlyElement(context.getTearDownFunctions()).run();
- assertThat(mainOutputValues, empty());
- }
- }
-
- @RunWith(JUnit4.class)
- public static class SplitTest {
- @Rule public final ExpectedException expected = ExpectedException.none();
- private IntervalWindow window1;
- private IntervalWindow window2;
- private IntervalWindow window3;
- private WindowedValue<String> currentElement;
- private OffsetRange currentRestriction;
- private Instant currentWatermarkEstimatorState;
- KV<Instant, Instant> watermarkAndState;
-
- private KV<WindowedValue, WindowedValue> createSplitAcrossWindows(
- List<BoundedWindow> primaryWindows, List<BoundedWindow> residualWindows) {
- return KV.of(
- primaryWindows.isEmpty()
- ? null
- : WindowedValues.of(
- KV.of(
- currentElement.getValue(),
- KV.of(currentRestriction, currentWatermarkEstimatorState)),
- currentElement.getTimestamp(),
- primaryWindows,
- currentElement.getPaneInfo()),
- residualWindows.isEmpty()
- ? null
- : WindowedValues.of(
- KV.of(
- currentElement.getValue(),
- KV.of(currentRestriction, currentWatermarkEstimatorState)),
- currentElement.getTimestamp(),
- residualWindows,
- currentElement.getPaneInfo()));
- }
-
- @Before
- public void setUp() {
- window1 = new IntervalWindow(Instant.ofEpochMilli(0), Instant.ofEpochMilli(10));
- window2 = new IntervalWindow(Instant.ofEpochMilli(10), Instant.ofEpochMilli(20));
- window3 = new IntervalWindow(Instant.ofEpochMilli(20), Instant.ofEpochMilli(30));
- currentElement =
- WindowedValues.of(
- "a",
- Instant.ofEpochMilli(57),
- ImmutableList.of(window1, window2, window3),
- PaneInfo.NO_FIRING);
- currentRestriction = new OffsetRange(0L, 100L);
- currentWatermarkEstimatorState = Instant.ofEpochMilli(21);
- watermarkAndState = KV.of(Instant.ofEpochMilli(42), Instant.ofEpochMilli(42));
- }
-
- private static final String PROCESS_TRANSFORM_ID = "processPTransformId";
- private static final String TRUNCATE_TRANSFORM_ID = "truncatePTransformId";
- private static final String PROCESS_INPUT_ID = "processInputId";
- private static final String TRUNCATE_INPUT_ID = "truncateInputId";
- private static final String TRUNCATE_OUTPUT_ID = "truncateOutputId";
-
- private HandlesSplits.SplitResult getProcessElementSplit(String transformId, String inputId) {
- return SplitResult.of(
- ImmutableList.of(
- BundleApplication.newBuilder()
- .setTransformId(transformId)
- .setInputId(inputId)
- .build()),
- ImmutableList.of(
- DelayedBundleApplication.newBuilder()
- .setApplication(
- BundleApplication.newBuilder()
- .setTransformId(transformId)
- .setInputId(inputId)
- .build())
- .setRequestedTimeDelay(Durations.fromMillis(1000L))
- .build()));
- }
-
- @Test
- public void testConstructSplitResultWithElementSplitFromDelegate() throws Exception {
- Coder fullInputCoder =
- getFullInputCoder(
- StringUtf8Coder.of(),
- OffsetRange.Coder.of(),
- InstantCoder.of(),
- IntervalWindow.getCoder());
- HandlesSplits.SplitResult elementSplit =
- getProcessElementSplit(PROCESS_TRANSFORM_ID, PROCESS_INPUT_ID);
- HandlesSplits.SplitResult result =
- FnApiDoFnRunner.constructSplitResult(
- null,
- elementSplit,
- fullInputCoder,
- null,
- null,
- TRUNCATE_TRANSFORM_ID,
- TRUNCATE_INPUT_ID,
- ImmutableList.of(TRUNCATE_OUTPUT_ID),
- null);
- assertEquals(elementSplit.getPrimaryRoots(), result.getPrimaryRoots());
- assertEquals(elementSplit.getResidualRoots(), result.getResidualRoots());
- }
-
- private HandlesSplits createSplitDelegate(
- double progress, double expectedFraction, SplitResult result) {
- return new HandlesSplits() {
- @Override
- public SplitResult trySplit(double fractionOfRemainder) {
- checkArgument(fractionOfRemainder == expectedFraction);
- return result;
- }
-
- @Override
- public double getProgress() {
- return progress;
- }
- };
- }
-
- @Test
- public void testTrySplitForTruncateCheckpointOnFirstWindow() throws Exception {
- List<BoundedWindow> windows = ImmutableList.copyOf(currentElement.getWindows());
- SplitResult splitResult =
- SplitResult.of(
- ImmutableList.of(BundleApplication.getDefaultInstance()),
- ImmutableList.of(DelayedBundleApplication.getDefaultInstance()));
- HandlesSplits splitDelegate = createSplitDelegate(0.3, 0.0, splitResult);
- SplitResultsWithStopIndex result =
- SplittableTruncateSizedRestrictionsDoFnRunner.computeSplitForTruncate(
- currentElement,
- currentRestriction,
- window1,
- windows,
- currentWatermarkEstimatorState,
- 0.0,
- splitDelegate,
- 0,
- 3);
- assertEquals(1, result.getNewWindowStopIndex());
- KV<WindowedValue, WindowedValue> expectedWindowSplit =
- createSplitAcrossWindows(ImmutableList.of(), ImmutableList.of(window2, window3));
- assertEquals(splitResult, result.getDownstreamSplit());
- assertNull(result.getWindowSplit().getPrimarySplitRoot());
- assertNull(result.getWindowSplit().getResidualSplitRoot());
- assertEquals(
- expectedWindowSplit.getKey(),
- result.getWindowSplit().getPrimaryInFullyProcessedWindowsRoot());
- assertEquals(
- expectedWindowSplit.getValue(),
- result.getWindowSplit().getResidualInUnprocessedWindowsRoot());
- }
-
- @Test
- public void testTrySplitForTruncateCheckpointOnFirstWindowAfterOneSplit() throws Exception {
- List<BoundedWindow> windows = ImmutableList.copyOf(currentElement.getWindows());
- SplitResult splitResult =
- SplitResult.of(
- ImmutableList.of(BundleApplication.getDefaultInstance()),
- ImmutableList.of(DelayedBundleApplication.getDefaultInstance()));
- HandlesSplits splitDelegate = createSplitDelegate(0.3, 0.0, splitResult);
- SplitResultsWithStopIndex result =
- SplittableTruncateSizedRestrictionsDoFnRunner.computeSplitForTruncate(
- currentElement,
- currentRestriction,
- window1,
- windows,
- currentWatermarkEstimatorState,
- 0.0,
- splitDelegate,
- 0,
- 2);
- assertEquals(1, result.getNewWindowStopIndex());
- KV<WindowedValue, WindowedValue> expectedWindowSplit =
- createSplitAcrossWindows(ImmutableList.of(), ImmutableList.of(window2));
- assertEquals(splitResult, result.getDownstreamSplit());
- assertNull(result.getWindowSplit().getPrimarySplitRoot());
- assertNull(result.getWindowSplit().getResidualSplitRoot());
- assertEquals(
- expectedWindowSplit.getKey(),
- result.getWindowSplit().getPrimaryInFullyProcessedWindowsRoot());
- assertEquals(
- expectedWindowSplit.getValue(),
- result.getWindowSplit().getResidualInUnprocessedWindowsRoot());
- }
-
- @Test
- public void testTrySplitForTruncateSplitOnFirstWindow() throws Exception {
- List<BoundedWindow> windows = ImmutableList.copyOf(currentElement.getWindows());
- SplitResult splitResult =
- SplitResult.of(
- ImmutableList.of(BundleApplication.getDefaultInstance()),
- ImmutableList.of(DelayedBundleApplication.getDefaultInstance()));
- HandlesSplits splitDelegate = createSplitDelegate(0.3, 0.54, splitResult);
- SplitResultsWithStopIndex result =
- SplittableTruncateSizedRestrictionsDoFnRunner.computeSplitForTruncate(
- currentElement,
- currentRestriction,
- window1,
- windows,
- currentWatermarkEstimatorState,
- 0.2,
- splitDelegate,
- 0,
- 3);
- assertEquals(1, result.getNewWindowStopIndex());
- KV<WindowedValue, WindowedValue> expectedWindowSplit =
- createSplitAcrossWindows(ImmutableList.of(), ImmutableList.of(window2, window3));
- assertEquals(splitResult, result.getDownstreamSplit());
- assertNull(result.getWindowSplit().getPrimarySplitRoot());
- assertNull(result.getWindowSplit().getResidualSplitRoot());
- assertEquals(
- expectedWindowSplit.getKey(),
- result.getWindowSplit().getPrimaryInFullyProcessedWindowsRoot());
- assertEquals(
- expectedWindowSplit.getValue(),
- result.getWindowSplit().getResidualInUnprocessedWindowsRoot());
- }
-
- @Test
- public void testTrySplitForTruncateSplitOnMiddleWindow() throws Exception {
- List<BoundedWindow> windows = ImmutableList.copyOf(currentElement.getWindows());
- SplitResult splitResult =
- SplitResult.of(
- ImmutableList.of(BundleApplication.getDefaultInstance()),
- ImmutableList.of(DelayedBundleApplication.getDefaultInstance()));
- HandlesSplits splitDelegate = createSplitDelegate(0.3, 0.34, splitResult);
- SplitResultsWithStopIndex result =
- SplittableTruncateSizedRestrictionsDoFnRunner.computeSplitForTruncate(
- currentElement,
- currentRestriction,
- window1,
- windows,
- currentWatermarkEstimatorState,
- 0.2,
- splitDelegate,
- 1,
- 3);
- assertEquals(2, result.getNewWindowStopIndex());
- KV<WindowedValue, WindowedValue> expectedWindowSplit =
- createSplitAcrossWindows(ImmutableList.of(window1), ImmutableList.of(window3));
- assertEquals(splitResult, result.getDownstreamSplit());
- assertNull(result.getWindowSplit().getPrimarySplitRoot());
- assertNull(result.getWindowSplit().getResidualSplitRoot());
- assertEquals(
- expectedWindowSplit.getKey(),
- result.getWindowSplit().getPrimaryInFullyProcessedWindowsRoot());
- assertEquals(
- expectedWindowSplit.getValue(),
- result.getWindowSplit().getResidualInUnprocessedWindowsRoot());
- }
-
- @Test
- public void testTrySplitForTruncateSplitOnLastWindow() throws Exception {
- List<BoundedWindow> windows = ImmutableList.copyOf(currentElement.getWindows());
- SplitResult splitResult =
- SplitResult.of(
- ImmutableList.of(BundleApplication.getDefaultInstance()),
- ImmutableList.of(DelayedBundleApplication.getDefaultInstance()));
- HandlesSplits splitDelegate = createSplitDelegate(0.3, 0.2, splitResult);
- SplitResultsWithStopIndex result =
- SplittableTruncateSizedRestrictionsDoFnRunner.computeSplitForTruncate(
- currentElement,
- currentRestriction,
- window1,
- windows,
- currentWatermarkEstimatorState,
- 0.2,
- splitDelegate,
- 2,
- 3);
- assertEquals(3, result.getNewWindowStopIndex());
- KV<WindowedValue, WindowedValue> expectedWindowSplit =
- createSplitAcrossWindows(ImmutableList.of(window1, window2), ImmutableList.of());
- assertEquals(splitResult, result.getDownstreamSplit());
- assertNull(result.getWindowSplit().getPrimarySplitRoot());
- assertNull(result.getWindowSplit().getResidualSplitRoot());
- assertEquals(
- expectedWindowSplit.getKey(),
- result.getWindowSplit().getPrimaryInFullyProcessedWindowsRoot());
- assertEquals(
- expectedWindowSplit.getValue(),
- result.getWindowSplit().getResidualInUnprocessedWindowsRoot());
- }
-
- @Test
- public void testTrySplitForTruncateSplitOnFirstWindowFallback() throws Exception {
- List<BoundedWindow> windows = ImmutableList.copyOf(currentElement.getWindows());
- SplitResult unusedSplitResult =
- SplitResult.of(
- ImmutableList.of(BundleApplication.getDefaultInstance()),
- ImmutableList.of(DelayedBundleApplication.getDefaultInstance()));
- HandlesSplits splitDelegate = createSplitDelegate(1.0, 0.0, unusedSplitResult);
- SplitResultsWithStopIndex result =
- SplittableTruncateSizedRestrictionsDoFnRunner.computeSplitForTruncate(
- currentElement,
- currentRestriction,
- window1,
- windows,
- currentWatermarkEstimatorState,
- 0.0,
- splitDelegate,
- 0,
- 3);
- assertEquals(1, result.getNewWindowStopIndex());
- KV<WindowedValue, WindowedValue> expectedWindowSplit =
- createSplitAcrossWindows(ImmutableList.of(window1), ImmutableList.of(window2, window3));
- assertNull(result.getDownstreamSplit());
- assertNull(result.getWindowSplit().getPrimarySplitRoot());
- assertNull(result.getWindowSplit().getResidualSplitRoot());
- assertEquals(
- expectedWindowSplit.getKey(),
- result.getWindowSplit().getPrimaryInFullyProcessedWindowsRoot());
- assertEquals(
- expectedWindowSplit.getValue(),
- result.getWindowSplit().getResidualInUnprocessedWindowsRoot());
- }
-
- @Test
- public void testTrySplitForTruncateSplitOnLastWindowWhenNoElementSplit() throws Exception {
- List<BoundedWindow> windows = ImmutableList.copyOf(currentElement.getWindows());
- HandlesSplits splitDelegate = createSplitDelegate(1.0, 0.0, null);
- SplitResultsWithStopIndex result =
- SplittableTruncateSizedRestrictionsDoFnRunner.computeSplitForTruncate(
- currentElement,
- currentRestriction,
- window1,
- windows,
- currentWatermarkEstimatorState,
- 0.0,
- splitDelegate,
- 2,
- 3);
- assertNull(result);
- }
-
- @Test
- public void testTrySplitForTruncateOnWindowBoundaryRoundUp() throws Exception {
- List<BoundedWindow> windows = ImmutableList.copyOf(currentElement.getWindows());
- SplitResult unusedSplitResult =
- SplitResult.of(
- ImmutableList.of(BundleApplication.getDefaultInstance()),
- ImmutableList.of(DelayedBundleApplication.getDefaultInstance()));
- HandlesSplits splitDelegate = createSplitDelegate(0.3, 0.0, unusedSplitResult);
- SplitResultsWithStopIndex result =
- SplittableTruncateSizedRestrictionsDoFnRunner.computeSplitForTruncate(
- currentElement,
- currentRestriction,
- window1,
- windows,
- currentWatermarkEstimatorState,
- 0.6,
- splitDelegate,
- 0,
- 3);
- assertEquals(2, result.getNewWindowStopIndex());
- KV<WindowedValue, WindowedValue> expectedWindowSplit =
- createSplitAcrossWindows(ImmutableList.of(window1, window2), ImmutableList.of(window3));
- assertNull(result.getDownstreamSplit());
- assertNull(result.getWindowSplit().getPrimarySplitRoot());
- assertNull(result.getWindowSplit().getResidualSplitRoot());
- assertEquals(
- expectedWindowSplit.getKey(),
- result.getWindowSplit().getPrimaryInFullyProcessedWindowsRoot());
- assertEquals(
- expectedWindowSplit.getValue(),
- result.getWindowSplit().getResidualInUnprocessedWindowsRoot());
- }
-
- @Test
- public void testTrySplitForTruncateOnWindowBoundaryRoundDown() throws Exception {
- List<BoundedWindow> windows = ImmutableList.copyOf(currentElement.getWindows());
- SplitResult unusedSplitResult =
- SplitResult.of(
- ImmutableList.of(BundleApplication.getDefaultInstance()),
- ImmutableList.of(DelayedBundleApplication.getDefaultInstance()));
- HandlesSplits splitDelegate = createSplitDelegate(0.3, 0.0, unusedSplitResult);
- SplitResultsWithStopIndex result =
- SplittableTruncateSizedRestrictionsDoFnRunner.computeSplitForTruncate(
- currentElement,
- currentRestriction,
- window1,
- windows,
- currentWatermarkEstimatorState,
- 0.3,
- splitDelegate,
- 0,
- 3);
- assertEquals(1, result.getNewWindowStopIndex());
- KV<WindowedValue, WindowedValue> expectedWindowSplit =
- createSplitAcrossWindows(ImmutableList.of(window1), ImmutableList.of(window2, window3));
- assertNull(result.getDownstreamSplit());
- assertNull(result.getWindowSplit().getPrimarySplitRoot());
- assertNull(result.getWindowSplit().getResidualSplitRoot());
- assertEquals(
- expectedWindowSplit.getKey(),
- result.getWindowSplit().getPrimaryInFullyProcessedWindowsRoot());
- assertEquals(
- expectedWindowSplit.getValue(),
- result.getWindowSplit().getResidualInUnprocessedWindowsRoot());
- }
-
- @Test
- public void testTrySplitForTruncateOnWindowBoundaryRoundDownOnLastWindow() throws Exception {
- List<BoundedWindow> windows = ImmutableList.copyOf(currentElement.getWindows());
- SplitResult unusedSplitResult =
- SplitResult.of(
- ImmutableList.of(BundleApplication.getDefaultInstance()),
- ImmutableList.of(DelayedBundleApplication.getDefaultInstance()));
- HandlesSplits splitDelegate = createSplitDelegate(0.3, 0.0, unusedSplitResult);
- SplitResultsWithStopIndex result =
- SplittableTruncateSizedRestrictionsDoFnRunner.computeSplitForTruncate(
- currentElement,
- currentRestriction,
- window1,
- windows,
- currentWatermarkEstimatorState,
- 0.6,
- splitDelegate,
- 0,
- 3);
- assertEquals(2, result.getNewWindowStopIndex());
- KV<WindowedValue, WindowedValue> expectedWindowSplit =
- createSplitAcrossWindows(ImmutableList.of(window1, window2), ImmutableList.of(window3));
- assertNull(result.getDownstreamSplit());
- assertNull(result.getWindowSplit().getPrimarySplitRoot());
- assertNull(result.getWindowSplit().getResidualSplitRoot());
- assertEquals(
- expectedWindowSplit.getKey(),
- result.getWindowSplit().getPrimaryInFullyProcessedWindowsRoot());
- assertEquals(
- expectedWindowSplit.getValue(),
- result.getWindowSplit().getResidualInUnprocessedWindowsRoot());
- }
-
- private Coder getFullInputCoder(
- Coder elementCoder, Coder restrictionCoder, Coder watermarkStateCoder, Coder windowCoder) {
- Coder inputCoder =
- KvCoder.of(
- KvCoder.of(elementCoder, KvCoder.of(restrictionCoder, watermarkStateCoder)),
- DoubleCoder.of());
- return WindowedValues.getFullCoder(inputCoder, windowCoder);
- }
- }
-}