@Override public PCollection<KV<K, OutputT>> expand(PCollection<KV<K, InputT>> input) { return input .apply(fewKeys ? GroupByKey.createWithFewKeys() : GroupByKey.create()) .apply( Combine.<K, InputT, OutputT>groupedValues(fn, fnDisplayData) .withSideInputs(sideInputs)); }
@Test public void testDisplayData() { GroupByKey<String, String> groupByKey = GroupByKey.create(); GroupByKey<String, String> groupByFewKeys = GroupByKey.createWithFewKeys(); DisplayData gbkDisplayData = DisplayData.from(groupByKey); DisplayData fewKeysDisplayData = DisplayData.from(groupByFewKeys); assertThat(gbkDisplayData.items(), empty()); assertThat(fewKeysDisplayData, hasDisplayItem("fewKeys", true)); }