Search in sources :

Example 1 with ByteKey

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

the class BigtableIOTest method testReadingWithKeyRange.

/**
   * Tests reading all rows using key ranges. Tests a prefix [), a suffix (], and a restricted
   * range [] and that some properties hold across them.
   */
@Test
public void testReadingWithKeyRange() throws Exception {
    final String table = "TEST-KEY-RANGE-TABLE";
    final int numRows = 1001;
    List<Row> testRows = makeTableData(table, numRows);
    ByteKey startKey = ByteKey.copyFrom("key000000100".getBytes());
    ByteKey endKey = ByteKey.copyFrom("key000000300".getBytes());
    service.setupSampleRowKeys(table, numRows / 10, "key000000100".length());
    // Test prefix: [beginning, startKey).
    final ByteKeyRange prefixRange = ByteKeyRange.ALL_KEYS.withEndKey(startKey);
    List<Row> prefixRows = filterToRange(testRows, prefixRange);
    runReadTest(defaultRead.withTableId(table).withKeyRange(prefixRange), prefixRows);
    // Test suffix: [startKey, end).
    final ByteKeyRange suffixRange = ByteKeyRange.ALL_KEYS.withStartKey(startKey);
    List<Row> suffixRows = filterToRange(testRows, suffixRange);
    runReadTest(defaultRead.withTableId(table).withKeyRange(suffixRange), suffixRows);
    // Test restricted range: [startKey, endKey).
    final ByteKeyRange middleRange = ByteKeyRange.of(startKey, endKey);
    List<Row> middleRows = filterToRange(testRows, middleRange);
    runReadTest(defaultRead.withTableId(table).withKeyRange(middleRange), middleRows);
    //////// Size and content sanity checks //////////
    // Prefix, suffix, middle should be non-trivial (non-zero,non-all).
    assertThat(prefixRows, allOf(hasSize(lessThan(numRows)), hasSize(greaterThan(0))));
    assertThat(suffixRows, allOf(hasSize(lessThan(numRows)), hasSize(greaterThan(0))));
    assertThat(middleRows, allOf(hasSize(lessThan(numRows)), hasSize(greaterThan(0))));
    // Prefix + suffix should be exactly all rows.
    List<Row> union = Lists.newArrayList(prefixRows);
    union.addAll(suffixRows);
    assertThat("prefix + suffix = total", union, containsInAnyOrder(testRows.toArray(new Row[] {})));
    // Suffix should contain the middle.
    assertThat(suffixRows, hasItems(middleRows.toArray(new Row[] {})));
}
Also used : ByteKey(org.apache.beam.sdk.io.range.ByteKey) ByteKeyRange(org.apache.beam.sdk.io.range.ByteKeyRange) ByteString(com.google.protobuf.ByteString) Row(com.google.bigtable.v2.Row) Test(org.junit.Test)

Example 2 with ByteKey

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

the class HBaseIOTest method testReadingWithKeyRange.

/**
     * Tests reading all rows using key ranges. Tests a prefix [), a suffix (], and a restricted
     * range [] and that some properties hold across them.
     */
@Test
public void testReadingWithKeyRange() throws Exception {
    final String table = "TEST-KEY-RANGE-TABLE";
    final int numRows = 1001;
    final byte[] startRow = "2".getBytes();
    final byte[] stopRow = "9".getBytes();
    final ByteKey startKey = ByteKey.copyFrom(startRow);
    createTable(table);
    writeData(table, numRows);
    // Test prefix: [beginning, startKey).
    final ByteKeyRange prefixRange = ByteKeyRange.ALL_KEYS.withEndKey(startKey);
    runReadTestLength(HBaseIO.read().withConfiguration(conf).withTableId(table).withKeyRange(prefixRange), 126);
    // Test suffix: [startKey, end).
    final ByteKeyRange suffixRange = ByteKeyRange.ALL_KEYS.withStartKey(startKey);
    runReadTestLength(HBaseIO.read().withConfiguration(conf).withTableId(table).withKeyRange(suffixRange), 875);
    // Test restricted range: [startKey, endKey).
    // This one tests the second signature of .withKeyRange
    runReadTestLength(HBaseIO.read().withConfiguration(conf).withTableId(table).withKeyRange(startRow, stopRow), 441);
}
Also used : ByteKey(org.apache.beam.sdk.io.range.ByteKey) ByteKeyRange(org.apache.beam.sdk.io.range.ByteKeyRange) Test(org.junit.Test)

Aggregations

ByteKey (org.apache.beam.sdk.io.range.ByteKey)2 ByteKeyRange (org.apache.beam.sdk.io.range.ByteKeyRange)2 Test (org.junit.Test)2 Row (com.google.bigtable.v2.Row)1 ByteString (com.google.protobuf.ByteString)1