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();
}
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);
}
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[] {})));
}
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));
}
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;
}
Aggregations