Fix nullness in extensions/avro

sdks/java/extensions/avro/src/main/java/org/apache/beam/sdk/extensions/avro/coders/AvroCoder.java

@@ -17,6 +17,8 @@
  */
 package org.apache.beam.sdk.extensions.avro.coders;
 
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull;
+
 import com.google.errorprone.annotations.FormatMethod;
 import com.google.errorprone.annotations.FormatString;
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
@@ -109,9 +111,6 @@
  *
  * @param <T> the type of elements handled by this coder
  */
-@SuppressWarnings({
-  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
 public class AvroCoder<T> extends CustomCoder<T> {
 
   private static final Cache<AvroCoderCacheKey, AvroCoder<?>> AVRO_CODER_CACHE =
@@ -137,7 +136,12 @@ public static <T> AvroCoder<T> specific(TypeDescriptor<T> type) {
    * suite for encoding and decoding.
    */
   public static <T> AvroCoder<T> specific(Class<T> type) {
-    return specific(type, new SpecificData(type.getClassLoader()).getSchema(type));
+    return specific(type, specificSchemaOf(type));
+  }
+
+  @SuppressWarnings("nullness") // SpecificData tolerates a null class loader but is unannotated
+  private static Schema specificSchemaOf(Class<?> type) {
+    return new SpecificData(type.getClassLoader()).getSchema(type);
   }
 
   /**
@@ -167,7 +171,12 @@ public static <T> AvroCoder<T> reflect(TypeDescriptor<T> type) {
    * suite for encoding and decoding.
    */
   public static <T> AvroCoder<T> reflect(Class<T> type) {
-    return reflect(type, new ReflectData(type.getClassLoader()).getSchema(type));
+    return reflect(type, reflectSchemaOf(type));
+  }
+
+  @SuppressWarnings("nullness") // ReflectData tolerates a null class loader but is unannotated
+  private static Schema reflectSchemaOf(Class<?> type) {
+    return new ReflectData(type.getClassLoader()).getSchema(type);
   }
 
   /**
@@ -395,10 +404,10 @@ public Schema get() {
 
   // writer and reader are unused but kept for serialization update compatibility.
   @SuppressWarnings("unused")
-  private final EmptyOnDeserializationThreadLocal<DatumWriter<T>> writer = null;
+  private final @Nullable EmptyOnDeserializationThreadLocal<DatumWriter<T>> writer = null;
 
   @SuppressWarnings("unused")
-  private final EmptyOnDeserializationThreadLocal<DatumReader<T>> reader = null;
+  private final @Nullable EmptyOnDeserializationThreadLocal<DatumReader<T>> reader = null;
 
   // datumReader and datumWriter are initialized in the constructor and
   // on deserialization (see readObject).
@@ -424,7 +433,8 @@ protected AvroCoder(AvroDatumFactory<T> datumFactory, Schema schema) {
     this.decoder = new EmptyOnDeserializationThreadLocal<>();
     this.encoder = new EmptyOnDeserializationThreadLocal<>();
 
-    initializeAvroDatumReaderAndWriter();
+    this.datumReader = datumFactory.apply(schema, schema);
+    this.datumWriter = datumFactory.apply(schema);
   }
 
   /** Returns the type this coder encodes/decodes. */
@@ -473,6 +483,11 @@ public T decode(InputStream inStream) throws IOException {
     BinaryDecoder decoderInstance = DECODER_FACTORY.directBinaryDecoder(inStream, decoder.get());
     // Save the potentially-new instance for later.
     decoder.set(decoderInstance);
+    return readWithoutReuse(decoderInstance);
+  }
+
+  @SuppressWarnings("nullness") // DatumReader.read accepts a null reuse but is unannotated
+  private T readWithoutReuse(BinaryDecoder decoderInstance) throws IOException {
     return datumReader.read(null, decoderInstance);
   }
 
@@ -808,10 +823,10 @@ private void checkMap(String context, TypeDescriptor<?> type, Schema schema) {
     }
 
     private void checkArray(String context, TypeDescriptor<?> type, Schema schema) {
-      TypeDescriptor<?> elementType = null;
+      TypeDescriptor<?> elementType;
       if (type.isArray()) {
         // The type is an array (with ordering)-> deterministic iff the element is deterministic.
-        elementType = type.getComponentType();
+        elementType = checkNotNull(type.getComponentType());
       } else if (isSubtypeOf(type, Collection.class)) {
         if (isSubtypeOf(type, List.class, SortedSet.class)) {
           // Ordered collection -> deterministic iff the element is deterministic
@@ -895,16 +910,12 @@ private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundE
       this.datumReader = cachedCoder.get().datumReader;
       this.datumWriter = cachedCoder.get().datumWriter;
     } else {
-      initializeAvroDatumReaderAndWriter();
+      Schema schema = this.schemaSupplier.get();
+      this.datumReader = this.datumFactory.apply(schema, schema);
+      this.datumWriter = this.datumFactory.apply(schema);
     }
   }
 
-  private void initializeAvroDatumReaderAndWriter() {
-    this.datumReader =
-        this.datumFactory.apply(this.schemaSupplier.get(), this.schemaSupplier.get());
-    this.datumWriter = this.datumFactory.apply(this.schemaSupplier.get());
-  }
-
   enum AvroCoderType {
     SPECIFIC,
     REFLECT;

sdks/java/extensions/avro/src/main/java/org/apache/beam/sdk/extensions/avro/io/AvroDatumFactory.java

@@ -34,9 +34,6 @@
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 /** Create {@link DatumReader} and {@link DatumWriter} for given schemas. */
-@SuppressWarnings({
-  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
 public abstract class AvroDatumFactory<T>
     implements AvroSource.DatumReaderFactory<T>, AvroSink.DatumWriterFactory<T> {
 
@@ -172,20 +169,25 @@ public ReflectDatumFactory(Class<T> type) {
 
     @Override
     public DatumReader<T> apply(Schema writer, Schema reader) {
-      ReflectData data = new ReflectData(type.getClassLoader());
+      ReflectData data = newReflectData(type);
       AvroUtils.addLogicalTypeConversions(data);
       return new ReflectDatumReader<>(writer, reader, data);
     }
 
     @Override
     public DatumWriter<T> apply(Schema writer) {
-      ReflectData data = new ReflectData(type.getClassLoader());
+      ReflectData data = newReflectData(type);
       AvroUtils.addLogicalTypeConversions(data);
       return new ReflectDatumWriter<>(writer, data);
     }
 
     public static <T> ReflectDatumFactory<T> of(Class<T> type) {
       return new ReflectDatumFactory<>(type);
     }
+
+    @SuppressWarnings("nullness") // ReflectData tolerates a null class loader but is unannotated
+    private static ReflectData newReflectData(Class<?> type) {
+      return new ReflectData(type.getClassLoader());
+    }
   }
 }