@Test
@Category(RunnableOnService.class)
public void withLambdaAndTypeDescriptorShouldSucceed() {
PCollection<String> values = p.apply(Create.of("1234", "3210", "0", "-12"));
PCollection<KV<Integer, String>> kvs =
values.apply(
WithKeys.of((String s) -> Integer.valueOf(s))
.withKeyType(TypeDescriptor.of(Integer.class)));
PAssert.that(kvs)
.containsInAnyOrder(
KV.of(1234, "1234"), KV.of(0, "0"), KV.of(-12, "-12"), KV.of(3210, "3210"));
p.run();
}
/** Helper function to generate KV test data. */
private List<KV<ByteString, ByteString>> generateTableData(int numRows) {
List<KV<ByteString, ByteString>> testData = new ArrayList<>(numRows);
for (int i = 0; i < numRows; ++i) {
ByteString key = ByteString.copyFromUtf8(String.format("key%09d", i));
ByteString value = ByteString.copyFromUtf8(String.format("value%09d", i));
testData.add(KV.of(key, value));
}
return testData;
}
private static <K, InputT, OutputT, W extends BoundedWindow>
List<WindowedValue<KV<K, OutputT>>> runGABW(
GroupAlsoByWindowsDoFnFactory<K, InputT, OutputT> gabwFactory,
WindowingStrategy<?, W> windowingStrategy,
K key,
Collection<WindowedValue<InputT>> values)
throws Exception {
final StateInternalsFactory<K> stateInternalsCache = new CachingStateInternalsFactory<K>();
List<WindowedValue<KV<K, OutputT>>> output =
processElement(
gabwFactory.forStrategy(windowingStrategy, stateInternalsCache),
KV.<K, Iterable<WindowedValue<InputT>>>of(key, values));
// Sanity check for corruption
for (WindowedValue<KV<K, OutputT>> value : output) {
assertThat(value.getValue().getKey(), equalTo(key));
}
return output;
}
/** 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;
}