sdks/java: re-enable nullness checks in WithKeys (#39506)

Re-enable nullness checking in WithKeys by removing Nullness suppression annotation.
Update keyType field and constructor parameter to be Nullable.
Use checkArgumentNotNull instead of Guava's checkNotNull.
Introduce a createWithConstantKey helper with scoped nullness suppression
to satisfy type variable bounds checks for constant key transform creation.

Author: kennknowles

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

@@ -17,9 +17,8 @@
  */
 package org.apache.beam.sdk.transforms;
 
-import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull;
+import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull;
 
-import javax.annotation.CheckForNull;
 import org.apache.beam.sdk.coders.CannotProvideCoderException;
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.coders.CoderRegistry;
@@ -56,9 +55,6 @@
  * @param <V> the type of the elements in the input {@code PCollection} and the values in the output
  *     {@code PCollection}
  */
-@SuppressWarnings({
-  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
 public class WithKeys<K, V> extends PTransform<PCollection<V>, PCollection<KV<K, V>>> {
   /**
    * Returns a {@code PTransform} that takes a {@code PCollection<V>} and returns a {@code
@@ -69,7 +65,7 @@ public class WithKeys<K, V> extends PTransform<PCollection<V>, PCollection<KV<K,
    * result {@link PTransform}.
    */
   public static <K, V> WithKeys<K, V> of(SerializableFunction<V, K> fn) {
-    checkNotNull(
+    checkArgumentNotNull(
         fn, "WithKeys constructed with null function. Did you mean WithKeys.of((Void) null)?");
     return new WithKeys<>(fn, null);
   }
@@ -79,20 +75,22 @@ public static <K, V> WithKeys<K, V> of(SerializableFunction<V, K> fn) {
    * PCollection<KV<K, V>>}, where each of the values in the input {@code PCollection} has been
    * paired with the given key.
    */
-  @SuppressWarnings("unchecked")
-  public static <K, V> WithKeys<K, V> of(final @Nullable K key) {
-    return new WithKeys<>(
-        value -> key,
-        (TypeDescriptor<K>)
-            (key == null ? TypeDescriptors.voids() : TypeDescriptor.of(key.getClass())));
+  public static <K, V> WithKeys<K, V> of(final K key) {
+    TypeDescriptor<K> keyType;
+    if (key == null) {
+      keyType = (TypeDescriptor<K>) TypeDescriptors.voids();
+    } else {
+      keyType = (TypeDescriptor<K>) TypeDescriptor.of(key.getClass());
+    }
+    return new WithKeys<>(value -> key, keyType); // createWithConstantKey(key, keyType);
   }
 
   /////////////////////////////////////////////////////////////////////////////
 
   private SerializableFunction<V, K> fn;
-  @CheckForNull private transient TypeDescriptor<K> keyType;
+  private transient @Nullable TypeDescriptor<K> keyType;
 
-  private WithKeys(SerializableFunction<V, K> fn, TypeDescriptor<K> keyType) {
+  private WithKeys(SerializableFunction<V, K> fn, @Nullable TypeDescriptor<K> keyType) {
     this.fn = fn;
     this.keyType = keyType;
   }
@@ -121,6 +119,7 @@ public KV<K, V> apply(V element) {
                   }
                 }));
 
+    TypeDescriptor<K> keyType = this.keyType;
     try {
       Coder<K> keyCoder;
       CoderRegistry coderRegistry = in.getPipeline().getCoderRegistry();