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);
}
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);
}
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);
}
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);
}
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();
}
Aggregations