Examples with ByteKeyRange org.apache.beam.sdk.io.range.ByteKeyRange used on opensource projects

Example 1 with ByteKeyRange

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

the class ByteKeyRangeTracker method trySplit.

@Override
public SplitResult<ByteKeyRange> trySplit(double fractionOfRemainder) {
    // No split on an empty range.
    if (NO_KEYS.equals(range) || (!range.getEndKey().isEmpty() && range.getStartKey().equals(range.getEndKey()))) {
        return null;
    }
    // There is no more remaining work after the entire range has been claimed.
    if (lastAttemptedKey != null && lastAttemptedKey.isEmpty()) {
        return null;
    }
    ByteKey unprocessedRangeStartKey = (lastAttemptedKey == null) ? range.getStartKey() : next(lastAttemptedKey);
    ByteKey endKey = range.getEndKey();
    // There is no more space for split.
    if (!endKey.isEmpty() && unprocessedRangeStartKey.compareTo(endKey) >= 0) {
        return null;
    }
    // trailing zeros when fraction is 0.
    if (fractionOfRemainder == 0.0) {
        // as the checkpoint.
        if (lastAttemptedKey == null) {
            // We update our current range to an interval that contains no elements.
            ByteKeyRange rval = range;
            range = range.getStartKey().isEmpty() ? NO_KEYS : ByteKeyRange.of(range.getStartKey(), range.getStartKey());
            return SplitResult.of(range, rval);
        } else {
            range = ByteKeyRange.of(range.getStartKey(), unprocessedRangeStartKey);
            return SplitResult.of(range, ByteKeyRange.of(unprocessedRangeStartKey, endKey));
        }
    }
    ByteKeyRange unprocessedRange = ByteKeyRange.of(unprocessedRangeStartKey, range.getEndKey());
    ByteKey splitPos;
    try {
        // The interpolateKey shouldn't return empty key. Please refer to {@link
        // ByteKeyRange#interpolateKey}.
        splitPos = unprocessedRange.interpolateKey(fractionOfRemainder);
        checkState(!splitPos.isEmpty());
    } catch (Exception e) {
        // There is no way to interpolate a key based on provided fraction.
        return null;
    }
    // Computed splitPos is out of current tracking restriction.
    if (!range.getEndKey().isEmpty() && splitPos.compareTo(range.getEndKey()) >= 0) {
        return null;
    }
    range = ByteKeyRange.of(range.getStartKey(), splitPos);
    return SplitResult.of(range, ByteKeyRange.of(splitPos, endKey));
}

Example 2 with ByteKeyRange

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

the class ByteKeyRangeTrackerTest method testCheckpointUnstartedForAllKeysRange.

@Test
public void testCheckpointUnstartedForAllKeysRange() throws Exception {
    ByteKeyRangeTracker tracker = ByteKeyRangeTracker.of(ByteKeyRange.ALL_KEYS);
    ByteKeyRange checkpoint = tracker.trySplit(0).getResidual();
    // We expect to get the original range back and that the current restriction
    // is effectively made empty.
    assertEquals(ByteKeyRange.ALL_KEYS, checkpoint);
    assertEquals(ByteKeyRangeTracker.NO_KEYS, tracker.currentRestriction());
    tracker.checkDone();
}

Example 3 with ByteKeyRange

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

the class ByteKeyRangeTrackerTest method testCheckpointJustStarted.

@Test
public void testCheckpointJustStarted() throws Exception {
    ByteKeyRangeTracker tracker = ByteKeyRangeTracker.of(ByteKeyRange.of(ByteKey.of(0x10), ByteKey.of(0xc0)));
    assertTrue(tracker.tryClaim(ByteKey.of(0x10)));
    ByteKeyRange checkpoint = tracker.trySplit(0).getResidual();
    assertEquals(ByteKeyRange.of(ByteKey.of(0x10), ByteKey.of(0x10, 0x00)), tracker.currentRestriction());
    assertEquals(ByteKeyRange.of(ByteKey.of(0x10, 0x00), ByteKey.of(0xc0)), checkpoint);
    tracker.checkDone();
}

Example 4 with ByteKeyRange

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

the class ByteKeyRangeTrackerTest method testTryClaim.

@Test
public void testTryClaim() throws Exception {
    ByteKeyRange range = ByteKeyRange.of(ByteKey.of(0x10), ByteKey.of(0xc0));
    ByteKeyRangeTracker tracker = ByteKeyRangeTracker.of(range);
    assertEquals(range, tracker.currentRestriction());
    assertTrue(tracker.tryClaim(ByteKey.of(0x10)));
    assertTrue(tracker.tryClaim(ByteKey.of(0x10, 0x00)));
    assertTrue(tracker.tryClaim(ByteKey.of(0x10, 0x00, 0x00)));
    assertTrue(tracker.tryClaim(ByteKey.of(0x50)));
    assertTrue(tracker.tryClaim(ByteKey.of(0x99)));
    assertFalse(tracker.tryClaim(ByteKey.of(0xc0)));
    tracker.checkDone();
}

Example 5 with ByteKeyRange

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

the class ByteKeyRangeTrackerTest method testCheckpointRegular.

@Test
public void testCheckpointRegular() throws Exception {
    ByteKeyRangeTracker tracker = ByteKeyRangeTracker.of(ByteKeyRange.of(ByteKey.of(0x10), ByteKey.of(0xc0)));
    assertTrue(tracker.tryClaim(ByteKey.of(0x50)));
    assertTrue(tracker.tryClaim(ByteKey.of(0x90)));
    ByteKeyRange checkpoint = tracker.trySplit(0).getResidual();
    assertEquals(ByteKeyRange.of(ByteKey.of(0x10), ByteKey.of(0x90, 0x00)), tracker.currentRestriction());
    assertEquals(ByteKeyRange.of(ByteKey.of(0x90, 0x00), ByteKey.of(0xc0)), checkpoint);
    tracker.checkDone();
}