@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; }
@Override public String apply(KV<String, Long> input) { return input.getKey() + ": " + input.getValue(); }