WithKeys extends PTransform<PCollection<V>, PCollection<KV<K, V>>> { public static <K, V> WithKeys<K, V> of(SerializableFunction<V, K> fn) { checkNotNull( fn, "WithKeys constructed with null function. Did you mean WithKeys.of((Void) null)?"); return new WithKeys<>(fn, null); } private  WithKeys(SerializableFunction<V, K> fn, TypeDescriptor<K> keyType); static WithKeys<K, V> of(SerializableFunction<V, K> fn); @SuppressWarnings("unchecked") static WithKeys<K, V> of(final @Nullable K key); WithKeys<K, V> withKeyType(TypeDescriptor<K> keyType); @Override PCollection<KV<K, V>> expand(PCollection<V> in);  }
@Test @Category(NeedsRunner.class) public void testExtractKeys() { PCollection<String> input = p.apply(Create.of(Arrays.asList(COLLECTION)).withCoder(StringUtf8Coder.of())); PCollection<KV<Integer, String>> output = input.apply(WithKeys.of(new LengthAsKey())); PAssert.that(output).containsInAnyOrder(WITH_KEYS); p.run(); }
@Test @Category(NeedsRunner.class) public void testConstantKeys() { PCollection<String> input = p.apply(Create.of(Arrays.asList(COLLECTION)).withCoder(StringUtf8Coder.of())); PCollection<KV<Integer, String>> output = input.apply(WithKeys.of(100)); PAssert.that(output).containsInAnyOrder(WITH_CONST_KEYS); p.run(); }
@Test @Category(NeedsRunner.class) public void testConstantVoidKeys() { PCollection<String> input = p.apply(Create.of(Arrays.asList(COLLECTION)).withCoder(StringUtf8Coder.of())); PCollection<KV<Void, String>> output = input.apply(WithKeys.of((Void) null)); PAssert.that(output).containsInAnyOrder(WITH_CONST_NULL_KEYS); p.run(); }
@Test public void testWithKeysGetName() { assertEquals("WithKeys", WithKeys.<Integer, String>of(100).getName()); }
@Test @Category(NeedsRunner.class) public void withLambdaAndNoTypeDescriptorShouldThrow() { PCollection<String> values = p.apply(Create.of("1234", "3210", "0", "-12")); values.apply("ApplyKeysWithWithKeys", WithKeys.of(Integer::valueOf)); thrown.expect(IllegalStateException.class); thrown.expectMessage("Unable to return a default Coder for ApplyKeysWithWithKeys"); p.run(); }