Search in sources :

Example 6 with Progress

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

the class GrowableOffsetRangeTrackerTest method testProgressBeforeStart.

@Test
public void testProgressBeforeStart() throws Exception {
    SimpleEstimator simpleEstimator = new SimpleEstimator();
    GrowableOffsetRangeTracker tracker = new GrowableOffsetRangeTracker(10L, simpleEstimator);
    simpleEstimator.setEstimateRangeEnd(20);
    Progress currentProcess = tracker.getProgress();
    assertEquals(0, currentProcess.getWorkCompleted(), 0.001);
    assertEquals(10, currentProcess.getWorkRemaining(), 0.001);
    simpleEstimator.setEstimateRangeEnd(15);
    currentProcess = tracker.getProgress();
    assertEquals(0, currentProcess.getWorkCompleted(), 0.001);
    assertEquals(5, currentProcess.getWorkRemaining(), 0.001);
    simpleEstimator.setEstimateRangeEnd(5);
    currentProcess = tracker.getProgress();
    assertEquals(0, currentProcess.getWorkCompleted(), 0.001);
    assertEquals(0, currentProcess.getWorkRemaining(), 0.001);
}
Also used : Progress(org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.Progress) Test(org.junit.Test)

Example 7 with Progress

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

the class GrowableOffsetRangeTrackerTest method testProgress.

@Test
public void testProgress() throws Exception {
    long start = 10L;
    SimpleEstimator simpleEstimator = new SimpleEstimator();
    GrowableOffsetRangeTracker tracker = new GrowableOffsetRangeTracker(start, simpleEstimator);
    long cur = 20L;
    assertTrue(tracker.tryClaim(cur));
    simpleEstimator.setEstimateRangeEnd(5L);
    Progress currentProgress = tracker.getProgress();
    assertEquals(cur - start, currentProgress.getWorkCompleted(), 0.001);
    assertEquals(0, currentProgress.getWorkRemaining(), 0.001);
    simpleEstimator.setEstimateRangeEnd(35L);
    currentProgress = tracker.getProgress();
    assertEquals(cur - start, currentProgress.getWorkCompleted(), 0.001);
    assertEquals(35L - cur, currentProgress.getWorkRemaining(), 0.001);
    simpleEstimator.setEstimateRangeEnd(25L);
    currentProgress = tracker.getProgress();
    assertEquals(cur - start, currentProgress.getWorkCompleted(), 0.001);
    assertEquals(25L - cur, currentProgress.getWorkRemaining(), 0.001);
    simpleEstimator.setEstimateRangeEnd(Long.MAX_VALUE);
    currentProgress = tracker.getProgress();
    assertEquals(cur - start, currentProgress.getWorkCompleted(), 0.001);
    assertEquals(Long.MAX_VALUE - cur, currentProgress.getWorkRemaining(), 0.001);
}
Also used : Progress(org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.Progress) Test(org.junit.Test)

Example 8 with Progress

use of org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.Progress 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);
}
Also used : Progress(org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.Progress) HasProgress(org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress) HasProgress(org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress) VisibleForTesting(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting)

Example 9 with Progress

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

the class OffsetByteRangeTrackerTest method progressTracked.

@Test
public void progressTracked() {
    assertTrue(tracker.tryClaim(OffsetByteProgress.of(Offset.of(123), 10)));
    assertTrue(tracker.tryClaim(OffsetByteProgress.of(Offset.of(124), 11)));
    when(unownedBacklogReader.computeMessageStats(Offset.of(125))).thenReturn(ComputeMessageStatsResponse.newBuilder().setMessageBytes(1000).build());
    Progress progress = tracker.getProgress();
    assertEquals(21, progress.getWorkCompleted(), .0001);
    assertEquals(1000, progress.getWorkRemaining(), .0001);
}
Also used : Progress(org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.Progress) Test(org.junit.Test)

Example 10 with Progress

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

the class PerSubscriptionPartitionSdfTest method getProgressUnboundedRangeDelegates.

@Test
public void getProgressUnboundedRangeDelegates() {
    Progress progress = Progress.from(0, 0.2);
    when(tracker.getProgress()).thenReturn(progress);
    assertTrue(DoubleMath.fuzzyEquals(progress.getWorkRemaining(), sdf.getSize(PARTITION, RESTRICTION), .0001));
    verify(tracker).getProgress();
}
Also used : Progress(org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.Progress) Test(org.junit.Test)

Aggregations

Progress (org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.Progress)14 Test (org.junit.Test)13 OffsetRange (org.apache.beam.sdk.io.range.OffsetRange)5 HasProgress (org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress)1 VisibleForTesting (org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting)1