(beam) branch master updated: sdks/java: re-enable nullness checks in WithKeys (#39506)

[radek]

This is an automated email from the ASF dual-hosted git repository.

stankiewicz pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git


The following commit(s) were added to refs/heads/master by this push:
     new ad1e278da3e sdks/java: re-enable nullness checks in WithKeys (#39506)
ad1e278da3e is described below

commit ad1e278da3e28d146f87cbbcb0f244eb2f5e306a
Author: Kenneth Knowles <k...@google.com>
AuthorDate: Tue Jul 28 06:44:14 2026 -0400

    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.
---
 .../org/apache/beam/sdk/transforms/WithKeys.java   | 27 +++++++++++-----------
 1 file changed, 13 insertions(+), 14 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 96072d8ec29..bdf9b07a4e6 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
@@ -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 @@ import org.checkerframework.checker.nullness.qual.Nullable;
  * @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 class WithKeys<K, V> extends 
PTransform<PCollection<V>, PCollection<KV<K,
    * 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 class WithKeys<K, V> extends 
PTransform<PCollection<V>, PCollection<KV<K,
                   }
                 }));
 
+    TypeDescriptor<K> keyType = this.keyType;
     try {
       Coder<K> keyCoder;
       CoderRegistry coderRegistry = in.getPipeline().getCoderRegistry();