Examples with BigtableSource org.apache.beam.sdk.io.gcp.bigtable.BigtableIO.BigtableSource used on opensource projects

Example 11 with BigtableSource

use of org.apache.beam.sdk.io.gcp.bigtable.BigtableIO.BigtableSource in project beam by apache.

the class BigtableIOTest method testReduceSplitsWithSomeNonAdjacentRanges.

/**
 * Tests reduce splits with few non adjacent ranges.
 */
@Test
public void testReduceSplitsWithSomeNonAdjacentRanges() throws Exception {
    final String table = "TEST-MANY-ROWS-SPLITS-TABLE";
    final int numRows = 10;
    final int numSamples = 10;
    final long bytesPerRow = 100L;
    final int maxSplit = 3;
    // Set up test table data and sample row keys for size estimation and splitting.
    makeTableData(table, numRows);
    service.setupSampleRowKeys(table, numSamples, bytesPerRow);
    // Construct few non contiguous key ranges [..1][1..2][3..4][4..5][6..7][8..9]
    List<ByteKeyRange> keyRanges = Arrays.asList(ByteKeyRange.of(ByteKey.EMPTY, createByteKey(1)), ByteKeyRange.of(createByteKey(1), createByteKey(2)), ByteKeyRange.of(createByteKey(3), createByteKey(4)), ByteKeyRange.of(createByteKey(4), createByteKey(5)), ByteKeyRange.of(createByteKey(6), createByteKey(7)), ByteKeyRange.of(createByteKey(8), createByteKey(9)));
    // Expected ranges after split and reduction by maxSplitCount is [..2][3..5][6..7][8..9]
    List<ByteKeyRange> expectedKeyRangesAfterReducedSplits = Arrays.asList(ByteKeyRange.of(ByteKey.EMPTY, createByteKey(2)), ByteKeyRange.of(createByteKey(3), createByteKey(5)), ByteKeyRange.of(createByteKey(6), createByteKey(7)), ByteKeyRange.of(createByteKey(8), createByteKey(9)));
    // Generate source and split it.
    BigtableSource source = new BigtableSource(config.withTableId(StaticValueProvider.of(table)), BigtableReadOptions.builder().setKeyRanges(StaticValueProvider.of(keyRanges)).build(), null);
    List<BigtableSource> splits = new ArrayList<>();
    for (ByteKeyRange range : keyRanges) {
        splits.add(source.withSingleRange(range));
    }
    List<BigtableSource> reducedSplits = source.reduceSplits(splits, null, maxSplit);
    List<ByteKeyRange> actualRangesAfterSplit = new ArrayList<>();
    for (BigtableSource splitSource : reducedSplits) {
        actualRangesAfterSplit.addAll(splitSource.getRanges());
    }
    assertAllSourcesHaveSingleRanges(reducedSplits);
    assertThat(actualRangesAfterSplit, IsIterableContainingInAnyOrder.containsInAnyOrder(expectedKeyRangesAfterReducedSplits.toArray()));
}