Examples with HasProgress org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress used on opensource projects

Example 1 with HasProgress

use of org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress in project beam by apache.

the class FnApiDoFnRunner method computeSplitForProcessOrTruncate.

@VisibleForTesting
static <WatermarkEstimatorStateT> SplitResultsWithStopIndex computeSplitForProcessOrTruncate(WindowedValue currentElement, Object currentRestriction, BoundedWindow currentWindow, List<BoundedWindow> windows, WatermarkEstimatorStateT currentWatermarkEstimatorState, double fractionOfRemainder, RestrictionTracker currentTracker, HandlesSplits splitDelegate, KV<Instant, WatermarkEstimatorStateT> watermarkAndState, int currentWindowIndex, int stopWindowIndex) {
    // We should only have currentTracker or splitDelegate.
    checkArgument((currentTracker != null) ^ (splitDelegate != null));
    // When we have currentTracker, the watermarkAndState should not be null.
    if (currentTracker != null) {
        checkNotNull(watermarkAndState);
    }
    WindowedSplitResult windowedSplitResult = null;
    HandlesSplits.SplitResult downstreamSplitResult = null;
    int newWindowStopIndex = stopWindowIndex;
    // on a future window.
    if (currentWindowIndex != stopWindowIndex - 1) {
        // Compute the fraction of the remainder relative to the scaled progress.
        Progress elementProgress;
        if (currentTracker != null) {
            if (currentTracker instanceof HasProgress) {
                elementProgress = ((HasProgress) currentTracker).getProgress();
            } else {
                elementProgress = Progress.from(0, 1);
            }
        } else {
            double elementCompleted = splitDelegate.getProgress();
            elementProgress = Progress.from(elementCompleted, 1 - elementCompleted);
        }
        Progress scaledProgress = scaleProgress(elementProgress, currentWindowIndex, stopWindowIndex);
        double scaledFractionOfRemainder = scaledProgress.getWorkRemaining() * fractionOfRemainder;
        // 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(currentElement, currentRestriction, currentWindow, windows, currentWatermarkEstimatorState, newWindowStopIndex, newWindowStopIndex, stopWindowIndex, null, watermarkAndState);
        } else {
            // Compute the element split with the scaled fraction.
            SplitResult<?> elementSplit = null;
            if (currentTracker != null) {
                elementSplit = currentTracker.trySplit(scaledFractionOfRemainder / elementProgress.getWorkRemaining());
            } else {
                downstreamSplitResult = splitDelegate.trySplit(scaledFractionOfRemainder);
            }
            newWindowStopIndex = currentWindowIndex + 1;
            int toIndex = (elementSplit == null && downstreamSplitResult == null) ? newWindowStopIndex : currentWindowIndex;
            windowedSplitResult = computeWindowSplitResult(currentElement, currentRestriction, currentWindow, windows, currentWatermarkEstimatorState, toIndex, newWindowStopIndex, stopWindowIndex, elementSplit, watermarkAndState);
        }
    } else {
        // We are on the last window then compute the element split with given fraction.
        SplitResult<?> elementSplitResult = null;
        newWindowStopIndex = stopWindowIndex;
        if (currentTracker != null) {
            elementSplitResult = currentTracker.trySplit(fractionOfRemainder);
        } else {
            downstreamSplitResult = splitDelegate.trySplit(fractionOfRemainder);
        }
        if (elementSplitResult == null && downstreamSplitResult == null) {
            return null;
        }
        windowedSplitResult = computeWindowSplitResult(currentElement, currentRestriction, currentWindow, windows, currentWatermarkEstimatorState, currentWindowIndex, stopWindowIndex, stopWindowIndex, elementSplitResult, watermarkAndState);
    }
    return SplitResultsWithStopIndex.of(windowedSplitResult, downstreamSplitResult, newWindowStopIndex);
}