From eb48b53dc5040a6e890e0e35b434aca3dbab5c5a Mon Sep 17 00:00:00 2001
From: Yi Hu <huuyyi@gmail.com>
Date: Fri, 24 Jul 2026 10:09:32 -0400
Subject: [PATCH] Revert "Merge pull request #39043: Improve WithKeys coder
 inference context"

This reverts commit a15b880c53c6d7f4bc6533ebce778881edb6e532.
---
 .../apache/beam/sdk/transforms/WithKeys.java  | 40 ++++++-------------
 .../beam/sdk/transforms/WithKeysTest.java     | 15 -------
 2 files changed, 13 insertions(+), 42 deletions(-)

diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/WithKeys.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/WithKeys.java
index 34af5811af18..96072d8ec29b 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/WithKeys.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/WithKeys.java
@@ -110,33 +110,27 @@ 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",
-            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)));
+            MapElements.via(
+                new SimpleFunction<V, KV<K, V>>() {
+                  @Override
+                  public KV<K, V> apply(V element) {
+                    return KV.of(fn.apply(element), element);
+                  }
+                }));
 
     try {
+      Coder<K> keyCoder;
       CoderRegistry coderRegistry = in.getPipeline().getCoderRegistry();
-      if (outputType == null) {
-        Coder<K> keyCoder = coderRegistry.getOutputCoder(fn, in.getCoder());
-        result.setCoder(KvCoder.of(keyCoder, in.getCoder()));
+      if (keyType == null) {
+        keyCoder = coderRegistry.getOutputCoder(fn, in.getCoder());
       } else {
-        result.setCoder(coderRegistry.getCoder(outputType, checkNotNull(inputType), in.getCoder()));
+        keyCoder = coderRegistry.getCoder(keyType);
       }
+      // TODO: Remove when we can set the coder inference context.
+      result.setCoder(KvCoder.of(keyCoder, in.getCoder()));
     } catch (CannotProvideCoderException exc) {
       if (keyType != null) {
         try {
@@ -157,12 +151,4 @@ 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);
-  }
 }
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/WithKeysTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/WithKeysTest.java
index e16b0fb77c5d..fd178f8e7649 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/WithKeysTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/WithKeysTest.java
@@ -172,21 +172,6 @@ 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() {