Search in sources :

Example 46 with OffsetRange

use of org.apache.beam.sdk.io.range.OffsetRange in project beam by apache.

the class GrowableOffsetRangeTrackerTest method testSplit.

@Test
public void testSplit() throws Exception {
    SimpleEstimator simpleEstimator = new SimpleEstimator();
    GrowableOffsetRangeTracker tracker = new GrowableOffsetRangeTracker(0L, simpleEstimator);
    assertTrue(tracker.tryClaim(0L));
    simpleEstimator.setEstimateRangeEnd(16L);
    // The split of infinite range results in one finite range and one infinite range.
    SplitResult res = tracker.trySplit(0.5);
    assertEquals(new OffsetRange(0, 8), res.getPrimary());
    assertEquals(new OffsetRange(0, 8), tracker.currentRestriction());
    assertEquals(new OffsetRange(8, Long.MAX_VALUE), res.getResidual());
    // After the first split, the tracker should track a finite range. Estimate offset should not
    // impact split.
    simpleEstimator.setEstimateRangeEnd(12L);
    res = tracker.trySplit(0.5);
    assertEquals(new OffsetRange(0, 4), res.getPrimary());
    assertEquals(new OffsetRange(0, 4), tracker.currentRestriction());
    assertEquals(new OffsetRange(4, 8), res.getResidual());
    assertFalse(tracker.tryClaim(4L));
    tracker.checkDone();
}
Also used : OffsetRange(org.apache.beam.sdk.io.range.OffsetRange) Test(org.junit.Test)

Example 47 with OffsetRange

use of org.apache.beam.sdk.io.range.OffsetRange in project beam by apache.

the class GrowableOffsetRangeTrackerTest method testLargeRange.

@Test
public void testLargeRange() throws Exception {
    SimpleEstimator simpleEstimator = new SimpleEstimator();
    GrowableOffsetRangeTracker tracker = new GrowableOffsetRangeTracker(Long.MIN_VALUE, simpleEstimator);
    simpleEstimator.setEstimateRangeEnd(Long.MAX_VALUE);
    Progress progress = tracker.getProgress();
    assertEquals(0, progress.getWorkCompleted(), 0.001);
    assertEquals(BigDecimal.valueOf(Long.MAX_VALUE).subtract(BigDecimal.valueOf(Long.MIN_VALUE), MathContext.DECIMAL128).doubleValue(), progress.getWorkRemaining(), 0.001);
    simpleEstimator.setEstimateRangeEnd(Long.MIN_VALUE);
    SplitResult res = tracker.trySplit(0);
    assertEquals(new OffsetRange(Long.MIN_VALUE, Long.MIN_VALUE), res.getPrimary());
    assertEquals(new OffsetRange(Long.MIN_VALUE, Long.MAX_VALUE), res.getResidual());
}
Also used : OffsetRange(org.apache.beam.sdk.io.range.OffsetRange) Progress(org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.Progress) Test(org.junit.Test)

Example 48 with OffsetRange

use of org.apache.beam.sdk.io.range.OffsetRange in project beam by apache.

the class OffsetRangeTrackerTest method testCheckDoneWhenNotDone.

@Test
public void testCheckDoneWhenNotDone() {
    OffsetRangeTracker tracker = new OffsetRangeTracker(new OffsetRange(100, 200));
    assertTrue(tracker.tryClaim(150L));
    assertTrue(tracker.tryClaim(175L));
    expected.expectMessage("Last attempted offset was 175 in range [100, 200), " + "claiming work in [176, 200) was not attempted");
    tracker.checkDone();
}
Also used : OffsetRange(org.apache.beam.sdk.io.range.OffsetRange) Test(org.junit.Test)

Example 49 with OffsetRange

use of org.apache.beam.sdk.io.range.OffsetRange in project beam by apache.

the class OffsetRangeTrackerTest method testDoneBeforeClaim.

@Test
public void testDoneBeforeClaim() throws Exception {
    expected.expectMessage("Last attempted offset should not be null. No work was claimed in non-empty range [100, 200)");
    OffsetRangeTracker tracker = new OffsetRangeTracker(new OffsetRange(100, 200));
    tracker.checkDone();
}
Also used : OffsetRange(org.apache.beam.sdk.io.range.OffsetRange) Test(org.junit.Test)

Example 50 with OffsetRange

use of org.apache.beam.sdk.io.range.OffsetRange in project beam by apache.

the class OffsetRangeTrackerTest method testCheckpointClaimedLast.

@Test
public void testCheckpointClaimedLast() throws Exception {
    OffsetRangeTracker tracker = new OffsetRangeTracker(new OffsetRange(100, 200));
    assertTrue(tracker.tryClaim(105L));
    assertTrue(tracker.tryClaim(110L));
    assertTrue(tracker.tryClaim(199L));
    SplitResult checkpoint = tracker.trySplit(0);
    assertEquals(new OffsetRange(100, 200), tracker.currentRestriction());
    assertNull(checkpoint);
    tracker.checkDone();
}
Also used : OffsetRange(org.apache.beam.sdk.io.range.OffsetRange) Test(org.junit.Test)

Aggregations

OffsetRange (org.apache.beam.sdk.io.range.OffsetRange)63 Test (org.junit.Test)53 Instant (org.joda.time.Instant)8 ArrayList (java.util.ArrayList)5 OffsetRangeTracker (org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker)5 Progress (org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.Progress)5 ProcessContinuation (org.apache.beam.sdk.transforms.DoFn.ProcessContinuation)4 PartitionMetadata (org.apache.beam.sdk.io.gcp.spanner.changestreams.model.PartitionMetadata)3 DoFn (org.apache.beam.sdk.transforms.DoFn)3 BigDecimal (java.math.BigDecimal)2 RestrictionTracker (org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker)2 Offset (com.google.cloud.pubsublite.Offset)1 SuppressFBWarnings (edu.umd.cs.findbugs.annotations.SuppressFBWarnings)1 Scope (io.opencensus.common.Scope)1 AttributeValue (io.opencensus.trace.AttributeValue)1 Tracer (io.opencensus.trace.Tracer)1 Tracing (io.opencensus.trace.Tracing)1 Serializable (java.io.Serializable)1 Map (java.util.Map)1 Optional (java.util.Optional)1