Examples with SourceWithKeys com.google.cloud.bigtable.beam.CloudBigtableIO.SourceWithKeys used on opensource projects

Example 1 with SourceWithKeys

use of com.google.cloud.bigtable.beam.CloudBigtableIO.SourceWithKeys in project java-bigtable-hbase by googleapis.

the class CloudBigtableIOTest method testSampleRowKeys.

@Test
public void testSampleRowKeys() throws Exception {
    List<KeyOffset> sampleRowKeys = new ArrayList<>();
    int count = (int) (AbstractSource.COUNT_MAX_SPLIT_COUNT * 3 - 5);
    byte[][] keys = Bytes.split("A".getBytes(), "Z".getBytes(), count - 2);
    long tabletSize = 2L * 1024L * 1024L * 1024L;
    long boundary = 0;
    for (byte[] currentKey : keys) {
        boundary += tabletSize;
        try {
            sampleRowKeys.add(KeyOffset.create(ByteString.copyFrom(currentKey), boundary));
        } catch (NoClassDefFoundError e) {
            // This could cause some problems for javadoc or cobertura because of the shading magic we
            // do.
            e.printStackTrace();
            return;
        }
    }
    Source source = (Source) CloudBigtableIO.read(scanConfig);
    source.setSampleRowKeys(sampleRowKeys);
    List<SourceWithKeys> splits = source.getSplits(20000);
    Collections.sort(splits, new Comparator<SourceWithKeys>() {

        @Override
        public int compare(SourceWithKeys o1, SourceWithKeys o2) {
            return ByteStringComparator.INSTANCE.compare(o1.getConfiguration().getStartRowByteString(), o2.getConfiguration().getStartRowByteString());
        }
    });
    Assert.assertTrue(splits.size() <= AbstractSource.COUNT_MAX_SPLIT_COUNT);
    Iterator<SourceWithKeys> iter = splits.iterator();
    SourceWithKeys last = iter.next();
    while (iter.hasNext()) {
        SourceWithKeys current = iter.next();
        Assert.assertTrue(Bytes.equals(current.getConfiguration().getZeroCopyStartRow(), last.getConfiguration().getZeroCopyStopRow()));
        // The last source will have a stop key of empty.
        if (iter.hasNext()) {
            Assert.assertTrue(Bytes.compareTo(current.getConfiguration().getZeroCopyStartRow(), current.getConfiguration().getZeroCopyStopRow()) < 0);
        }
        Assert.assertTrue(current.getEstimatedSize() >= tabletSize);
        last = current;
    }
// check first and last
}