Merge pull request #39043: Improve WithKeys coder inference context

Author: adityaRaj369

sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/WithKeys.java

@@ -110,27 +110,33 @@ public WithKeys<K, V> withKeyType(TypeDescriptor<K> keyType) {
 
   @Override
   public PCollection<KV<K, V>> expand(PCollection<V> in) {
+    SerializableFunction<V, K> localFn = fn;
+    TypeDescriptor<V> inputType = in.getTypeDescriptor();
+    TypeDescriptor<KV<K, V>> outputType = getOutputTypeDescriptor(inputType);
     PCollection<KV<K, V>> result =
         in.apply(
             "AddKeys",
-            MapElements.via(
-                new SimpleFunction<V, KV<K, V>>() {
-                  @Override
-                  public KV<K, V> apply(V element) {
-                    return KV.of(fn.apply(element), element);
-                  }
-                }));
+            outputType == null
+                ? MapElements.via(
+                    new SimpleFunction<V, KV<K, V>>() {
+                      @Override
+                      public KV<K, V> apply(V element) {
+                        return KV.of(localFn.apply(element), element);
+                      }
+                    })
+                : MapElements.into(outputType)
+                    .via(
+                        (SerializableFunction<V, KV<K, V>>)
+                            element -> KV.of(localFn.apply(element), element)));
 
     try {
-      Coder<K> keyCoder;
       CoderRegistry coderRegistry = in.getPipeline().getCoderRegistry();
-      if (keyType == null) {
-        keyCoder = coderRegistry.getOutputCoder(fn, in.getCoder());
+      if (outputType == null) {
+        Coder<K> keyCoder = coderRegistry.getOutputCoder(fn, in.getCoder());
+        result.setCoder(KvCoder.of(keyCoder, in.getCoder()));
       } else {
-        keyCoder = coderRegistry.getCoder(keyType);
+        result.setCoder(coderRegistry.getCoder(outputType, checkNotNull(inputType), in.getCoder()));
       }
-      // TODO: Remove when we can set the coder inference context.
-      result.setCoder(KvCoder.of(keyCoder, in.getCoder()));
     } catch (CannotProvideCoderException exc) {
       if (keyType != null) {
         try {
@@ -151,4 +157,12 @@ public KV<K, V> apply(V element) {
 
     return result;
   }
+
+  private @Nullable TypeDescriptor<KV<K, V>> getOutputTypeDescriptor(
+      @Nullable TypeDescriptor<V> inputType) {
+    if (keyType == null || inputType == null) {
+      return null;
+    }
+    return TypeDescriptors.kvs(keyType, inputType);
+  }
 }

sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/WithKeysTest.java

@@ -172,6 +172,21 @@ public void withLambdaAndParameterizedTypeDescriptorShouldSucceed() {
     p.run();
   }
 
+  @Test
+  public void withKeyTypeShouldSetOutputTypeDescriptorFromInputType() {
+    PCollection<String> values =
+        p.apply(Create.of("1234", "3210").withType(TypeDescriptors.strings()));
+
+    PCollection<KV<Integer, String>> kvs =
+        values.apply(
+            WithKeys.of((SerializableFunction<String, Integer>) Integer::valueOf)
+                .withKeyType(TypeDescriptors.integers()));
+
+    assertEquals(
+        TypeDescriptors.kvs(TypeDescriptors.integers(), TypeDescriptors.strings()),
+        kvs.getTypeDescriptor());
+  }
+
   @Test
   @Category(NeedsRunner.class)
   public void withLambdaAndNoTypeDescriptorShouldThrow() {