Search in sources :

Example 6 with Row

use of com.google.bigtable.v2.Row in project beam by apache.

the class BigtableIOTest method testGetSplitPointsConsumed.

@Test
public void testGetSplitPointsConsumed() throws Exception {
    final String table = "TEST-TABLE";
    final int numRows = 100;
    int splitPointsConsumed = 0;
    makeTableData(table, numRows);
    BigtableSource source = new BigtableSource(serviceFactory, table, null, ByteKeyRange.ALL_KEYS, null);
    BoundedReader<Row> reader = source.createReader(TestPipeline.testingPipelineOptions());
    reader.start();
    // Started, 0 split points consumed
    assertEquals("splitPointsConsumed starting", splitPointsConsumed, reader.getSplitPointsConsumed());
    // Split points consumed increases for each row read
    while (reader.advance()) {
        assertEquals("splitPointsConsumed advancing", ++splitPointsConsumed, reader.getSplitPointsConsumed());
    }
    // Reader marked as done, 100 split points consumed
    assertEquals("splitPointsConsumed done", numRows, reader.getSplitPointsConsumed());
    reader.close();
}
Also used : ByteString(com.google.protobuf.ByteString) Row(com.google.bigtable.v2.Row) BigtableSource(org.apache.beam.sdk.io.gcp.bigtable.BigtableIO.BigtableSource) Test(org.junit.Test)

Example 7 with Row

use of com.google.bigtable.v2.Row in project beam by apache.

the class BigtableIOTest method makeWrite.

/** Helper function to make a single row mutation to be written. */
private static KV<ByteString, Iterable<Mutation>> makeWrite(String key, String value) {
    ByteString rowKey = ByteString.copyFromUtf8(key);
    Iterable<Mutation> mutations = ImmutableList.of(Mutation.newBuilder().setSetCell(SetCell.newBuilder().setValue(ByteString.copyFromUtf8(value))).build());
    return KV.of(rowKey, mutations);
}
Also used : ByteString(com.google.protobuf.ByteString) Mutation(com.google.bigtable.v2.Mutation)

Example 8 with Row

use of com.google.bigtable.v2.Row 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 9 with Row

use of com.google.bigtable.v2.Row in project beam by apache.

the class BigtableIOTest method testReading.

/** Tests reading all rows from a table. */
@Test
public void testReading() throws Exception {
    final String table = "TEST-MANY-ROWS-TABLE";
    final int numRows = 1001;
    List<Row> testRows = makeTableData(table, numRows);
    service.setupSampleRowKeys(table, 3, 1000L);
    runReadTest(defaultRead.withTableId(table), testRows);
    logged.verifyInfo(String.format("Closing reader after reading %d records.", numRows / 3));
}
Also used : ByteString(com.google.protobuf.ByteString) Row(com.google.bigtable.v2.Row) Test(org.junit.Test)

Example 10 with Row

use of com.google.bigtable.v2.Row in project beam by apache.

the class BigtableWriteIT method getTableData.

/** Helper function to get a table's data. */
private List<KV<ByteString, ByteString>> getTableData(String tableName) throws IOException {
    // Add empty range to avoid TARGET_NOT_SET error
    RowRange range = RowRange.newBuilder().setStartKeyClosed(ByteString.EMPTY).setEndKeyOpen(ByteString.EMPTY).build();
    RowSet rowSet = RowSet.newBuilder().addRowRanges(range).build();
    ReadRowsRequest.Builder readRowsRequestBuilder = ReadRowsRequest.newBuilder().setTableName(tableName).setRows(rowSet);
    ResultScanner<Row> scanner = session.getDataClient().readRows(readRowsRequestBuilder.build());
    Row currentRow;
    List<KV<ByteString, ByteString>> tableData = new ArrayList<>();
    while ((currentRow = scanner.next()) != null) {
        ByteString key = currentRow.getKey();
        ByteString value = currentRow.getFamilies(0).getColumns(0).getCells(0).getValue();
        tableData.add(KV.of(key, value));
    }
    scanner.close();
    return tableData;
}
Also used : RowRange(com.google.bigtable.v2.RowRange) ByteString(com.google.protobuf.ByteString) RowSet(com.google.bigtable.v2.RowSet) ArrayList(java.util.ArrayList) ReadRowsRequest(com.google.bigtable.v2.ReadRowsRequest) Row(com.google.bigtable.v2.Row) KV(org.apache.beam.sdk.values.KV)

Aggregations

ByteString (com.google.protobuf.ByteString)10 Row (com.google.bigtable.v2.Row)7 Test (org.junit.Test)7 RowFilter (com.google.bigtable.v2.RowFilter)3 ByteString (com.google.bigtable.repackaged.com.google.protobuf.ByteString)2 Column (com.google.bigtable.v1.Column)2 Family (com.google.bigtable.v1.Family)2 DeleteFromRow (com.google.bigtable.v1.Mutation.DeleteFromRow)2 ReadRowsRequest (com.google.bigtable.v1.ReadRowsRequest)2 Row (com.google.bigtable.v1.Row)2 RowFilter (com.google.bigtable.v1.RowFilter)2 Builder (com.google.bigtable.v1.RowFilter.Chain.Builder)2 ByteArrayByteIterator (com.yahoo.ycsb.ByteArrayByteIterator)2 ArrayList (java.util.ArrayList)2 ExecutionException (java.util.concurrent.ExecutionException)2 BigtableSource (org.apache.beam.sdk.io.gcp.bigtable.BigtableIO.BigtableSource)2 RowRange (com.google.bigtable.v1.RowRange)1 Mutation (com.google.bigtable.v2.Mutation)1 ReadRowsRequest (com.google.bigtable.v2.ReadRowsRequest)1 RowRange (com.google.bigtable.v2.RowRange)1