Repository: apache/beam
Pull request #34873: 34749 added cache for avro coder to reduce memory footprint
Author: login: wollowizard
Labels: name: java
name: extensions
name: Next Action: Reviewers
Created: 2025-05-06T21:19:55Z | Merged: 2025-05-12T08:15:08Z

This addresses issue https://github.com/apache/beam/issues/34749.

Results on heap allocation:

**AvroDatumFactory**
before: 32%
after: 0.036%

**AvroCoder**
before: 69%
after: 11%


As we thought, this improves not only user code, but also sdk performance.
Example taken from another pipeline:
**org.apache.beam.runners.dataflow.worker.StreamingModeExecutionContext.flushState()**
before: 43%
after: 10%


**Before**

<img width="600" alt="image" src="https://github.com/user-attachments/assets/b121e360-79c9-4fbd-b9d9-079ed9ad50c3" />
<img width="600" alt="image" src="https://github.com/user-attachments/assets/3a0e3651-c5f3-4d0b-a2ef-87ec61733bfe" />
<img width="600" alt="image" src="https://github.com/user-attachments/assets/6173ea0e-a9b6-4a23-9653-d62df14ef61e" />





**After**

<img width="600" alt="image" src="https://github.com/user-attachments/assets/a33c71f6-af23-481e-a721-5fb272161e67" />
<img width="600" alt="image" src="https://github.com/user-attachments/assets/31ee28a8-637c-4bd6-a95e-7eca0b19040b" />
<img width="600" alt="image" src="https://github.com/user-attachments/assets/ae1516dc-d23a-4857-ae09-e2c28e4977a0" />


See also PR https://github.com/apache/beam/pull/34750 for more details



Diff:
diff --git a/sdks/java/extensions/avro/src/main/java/org/apache/beam/sdk/extensions/avro/coders/AvroCoder.java b/sdks/java/extensions/avro/src/main/java/org/apache/beam/sdk/extensions/avro/coders/AvroCoder.java
index 826f209283fc..f943c098119f 100644
--- a/sdks/java/extensions/avro/src/main/java/org/apache/beam/sdk/extensions/avro/coders/AvroCoder.java
+++ b/sdks/java/extensions/avro/src/main/java/org/apache/beam/sdk/extensions/avro/coders/AvroCoder.java
@@ -22,6 +22,7 @@
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
 import java.io.IOException;
 import java.io.InputStream;
+import java.io.ObjectInputStream;
 import java.io.OutputStream;
 import java.io.Serializable;
 import java.lang.reflect.Field;
@@ -31,9 +32,13 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
 import java.util.SortedMap;
 import java.util.SortedSet;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutionException;
+import java.util.function.Supplier;
 import org.apache.avro.AvroRuntimeException;
 import org.apache.avro.Schema;
 import org.apache.avro.generic.GenericDatumWriter;
@@ -63,7 +68,8 @@
 import org.apache.beam.sdk.extensions.avro.io.AvroDatumFactory;
 import org.apache.beam.sdk.util.EmptyOnDeserializationThreadLocal;
 import org.apache.beam.sdk.values.TypeDescriptor;
-import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Supplier;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.Cache;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheBuilder;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 /**
@@ -108,6 +114,9 @@
 })
 public class AvroCoder<T> extends CustomCoder<T> {
 
+  private static final Cache<AvroCoderCacheKey, AvroCoder<?>> AVRO_CODER_CACHE =
+      CacheBuilder.newBuilder().weakValues().build();
+
   /**
    * Returns an {@link AvroCoder} instance for the Avro schema. The implicit type is GenericRecord.
    */
@@ -138,7 +147,11 @@ public static <T> AvroCoder<T> specific(Class<T> type) {
    * <p>The schema must correspond to the type provided.
    */
   public static <T> AvroCoder<T> specific(Class<T> type, Schema schema) {
-    return new AvroCoder<>(AvroDatumFactory.specific(type), schema);
+    return fromCacheOrCreate(
+        type,
+        schema,
+        AvroCoderType.SPECIFIC,
+        () -> new AvroCoder<>(AvroDatumFactory.specific(type), schema));
   }
 
   /**
@@ -164,7 +177,11 @@ public static <T> AvroCoder<T> reflect(Class<T> type) {
    * <p>The schema must correspond to the type provided.
    */
   public static <T> AvroCoder<T> reflect(Class<T> type, Schema schema) {
-    return new AvroCoder<>(AvroDatumFactory.reflect(type), schema);
+    return fromCacheOrCreate(
+        type,
+        schema,
+        AvroCoderType.REFLECT,
+        () -> new AvroCoder<>(AvroDatumFactory.reflect(type), schema));
   }
 
   /**
@@ -242,7 +259,28 @@ public static <T> AvroCoder<T> of(Class<T> type, Schema schema) {
    * @param <T> the element type
    */
   public static <T> AvroCoder<T> of(AvroDatumFactory<T> datumFactory, Schema schema) {
-    return new AvroCoder<>(datumFactory, schema);
+    Optional<AvroCoderType> avroCoderTypeOptional = AvroCoderType.fromDatumFactory(datumFactory);
+    if (!avroCoderTypeOptional.isPresent()) {
+      return new AvroCoder<>(datumFactory, schema);
+    }
+
+    Class<T> type = datumFactory.getType();
+    return fromCacheOrCreate(
+        type, schema, avroCoderTypeOptional.get(), () -> new AvroCoder<>(datumFactory, schema));
+  }
+
+  private static <T> AvroCoder<T> fromCacheOrCreate(
+      Class<T> type,
+      Schema schema,
+      AvroCoderType avroCoderType,
+      Callable<AvroCoder<T>> avroCoderCreator) {
+    try {
+      return (AvroCoder<T>)
+          AVRO_CODER_CACHE.get(
+              new AvroCoderCacheKey(type, schema, avroCoderType), avroCoderCreator);
+    } catch (ExecutionException e) {
+      throw new RuntimeException(e);
+    }
   }
 
   /**
@@ -354,8 +392,18 @@ public Schema get() {
   // an inner coder.
   private final EmptyOnDeserializationThreadLocal<BinaryDecoder> decoder;
   private final EmptyOnDeserializationThreadLocal<BinaryEncoder> encoder;
-  private final EmptyOnDeserializationThreadLocal<DatumWriter<T>> writer;
-  private final EmptyOnDeserializationThreadLocal<DatumReader<T>> reader;
+
+  // writer and reader are unused but kept for serialization update compatibility.
+  @SuppressWarnings("unused")
+  private final EmptyOnDeserializationThreadLocal<DatumWriter<T>> writer = null;
+
+  @SuppressWarnings("unused")
+  private final EmptyOnDeserializationThreadLocal<DatumReader<T>> reader = null;
+
+  // datumReader and datumWriter are initialized in the constructor and
+  // on deserialization (see readObject).
+  private transient DatumWriter<T> datumWriter;
+  private transient DatumReader<T> datumReader;
 
   protected AvroCoder(Class<T> type, Schema schema) {
     this(type, schema, false);
@@ -376,27 +424,7 @@ protected AvroCoder(AvroDatumFactory<T> datumFactory, Schema schema) {
     this.decoder = new EmptyOnDeserializationThreadLocal<>();
     this.encoder = new EmptyOnDeserializationThreadLocal<>();
 
-    // Reader and writer are allocated once per thread per Coder
-    this.reader =
-        new EmptyOnDeserializationThreadLocal<DatumReader<T>>() {
-          private final AvroCoder<T> myCoder = AvroCoder.this;
-
-          @Override
-          public DatumReader<T> initialValue() {
-            Schema schema = myCoder.getSchema();
-            return myCoder.datumFactory.apply(schema, schema);
-          }
-        };
-
-    this.writer =
-        new EmptyOnDeserializationThreadLocal<DatumWriter<T>>() {
-          private final AvroCoder<T> myCoder = AvroCoder.this;
-
-          @Override
-          public DatumWriter<T> initialValue() {
-            return myCoder.datumFactory.apply(myCoder.getSchema());
-          }
-        };
+    initializeAvroDatumReaderAndWriter();
   }
 
   /** Returns the type this coder encodes/decodes. */
@@ -409,6 +437,16 @@ public AvroDatumFactory<T> getDatumFactory() {
     return datumFactory;
   }
 
+  /** Returns the DatumWriter used for encoding. */
+  public DatumWriter<T> getDatumWriter() {
+    return datumWriter;
+  }
+
+  /** Returns the DatumReader used for decoding. */
+  public DatumReader<T> getDatumReader() {
+    return datumReader;
+  }
+
   /**
    * @return true if internal {@link AvroDatumFactory} is an instance of {@link
    *     AvroDatumFactory.ReflectDatumFactory}
@@ -425,7 +463,7 @@ public void encode(T value, OutputStream outStream) throws IOException {
     BinaryEncoder encoderInstance = ENCODER_FACTORY.directBinaryEncoder(outStream, encoder.get());
     // Save the potentially-new instance for reuse later.
     encoder.set(encoderInstance);
-    writer.get().write(value, encoderInstance);
+    datumWriter.write(value, encoderInstance);
     // Direct binary encoder does not buffer any data and need not be flushed.
   }
 
@@ -435,7 +473,7 @@ 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 reader.get().read(null, decoderInstance);
+    return datumReader.read(null, decoderInstance);
   }
 
   /**
@@ -840,4 +878,80 @@ public boolean equals(@Nullable Object other) {
   public int hashCode() {
     return Objects.hash(getClass(), typeDescriptor, datumFactory, schemaSupplier.get());
   }
+
+  private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundException {
+    in.defaultReadObject();
+    Optional<AvroCoder<T>> cachedCoder =
+        AvroCoderType.fromDatumFactory(this.datumFactory)
+            .map(
+                avroCoderType ->
+                    new AvroCoderCacheKey(
+                        this.typeDescriptor.getRawType(), this.schemaSupplier.get(), avroCoderType))
+            .map(
+                avroCoderCacheKey ->
+                    (AvroCoder<T>) AVRO_CODER_CACHE.getIfPresent(avroCoderCacheKey));
+
+    if (cachedCoder.isPresent()) {
+      this.datumReader = cachedCoder.get().datumReader;
+      this.datumWriter = cachedCoder.get().datumWriter;
+    } else {
+      initializeAvroDatumReaderAndWriter();
+    }
+  }
+
+  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;
+
+    public static Optional<AvroCoderType> fromDatumFactory(AvroDatumFactory<?> datumFactory) {
+      if (AvroDatumFactory.ReflectDatumFactory.class.equals(datumFactory.getClass())) {
+        return Optional.of(AvroCoderType.REFLECT);
+      } else if (AvroDatumFactory.SpecificDatumFactory.class.equals(datumFactory.getClass())) {
+        return Optional.of(AvroCoderType.SPECIFIC);
+      }
+      return Optional.empty();
+    }
+  }
+
+  static class AvroCoderCacheKey implements Serializable {
+    private final Class<?> type;
+    private final Schema schema;
+    private final AvroCoderType avroCoderType;
+
+    AvroCoderCacheKey(Class<?> type, Schema schema, AvroCoderType avroCoderType) {
+      this.type = type;
+      this.schema = schema;
+      this.avroCoderType = avroCoderType;
+    }
+
+    public static AvroCoderCacheKey specific(Class<?> type, Schema schema) {
+      return new AvroCoderCacheKey(type, schema, AvroCoderType.SPECIFIC);
+    }
+
+    public static AvroCoderCacheKey reflect(Class<?> type, Schema schema) {
+      return new AvroCoderCacheKey(type, schema, AvroCoderType.REFLECT);
+    }
+
+    @Override
+    public boolean equals(@Nullable Object o) {
+      if (o == null || getClass() != o.getClass()) {
+        return false;
+      }
+      AvroCoderCacheKey that = (AvroCoderCacheKey) o;
+      return Objects.equals(type, that.type)
+          && Objects.equals(schema, that.schema)
+          && Objects.equals(avroCoderType, that.avroCoderType);
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hash(type, schema, avroCoderType);
+    }
+  }
 }
diff --git a/sdks/java/extensions/avro/src/main/java/org/apache/beam/sdk/extensions/avro/coders/AvroGenericCoder.java b/sdks/java/extensions/avro/src/main/java/org/apache/beam/sdk/extensions/avro/coders/AvroGenericCoder.java
index 84acf75ceed6..74a244ad4099 100644
--- a/sdks/java/extensions/avro/src/main/java/org/apache/beam/sdk/extensions/avro/coders/AvroGenericCoder.java
+++ b/sdks/java/extensions/avro/src/main/java/org/apache/beam/sdk/extensions/avro/coders/AvroGenericCoder.java
@@ -17,17 +17,27 @@
  */
 package org.apache.beam.sdk.extensions.avro.coders;
 
+import java.util.concurrent.ExecutionException;
 import org.apache.avro.Schema;
 import org.apache.avro.generic.GenericRecord;
 import org.apache.beam.sdk.extensions.avro.io.AvroDatumFactory;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.Cache;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheBuilder;
 
 /** AvroCoder specialisation for GenericRecord, needed for cross-language transforms. */
 public class AvroGenericCoder extends AvroCoder<GenericRecord> {
+  private static final Cache<Schema, AvroGenericCoder> AVRO_GENERIC_CODER_CACHE =
+      CacheBuilder.newBuilder().weakValues().build();
+
   AvroGenericCoder(Schema schema) {
     super(AvroDatumFactory.GenericDatumFactory.INSTANCE, schema);
   }
 
   public static AvroGenericCoder of(Schema schema) {
-    return new AvroGenericCoder(schema);
+    try {
+      return AVRO_GENERIC_CODER_CACHE.get(schema, () -> new AvroGenericCoder(schema));
+    } catch (ExecutionException e) {
+      throw new RuntimeException(e);
+    }
   }
 }
diff --git a/sdks/java/extensions/avro/src/test/java/org/apache/beam/sdk/extensions/avro/coders/AvroCoderTest.java b/sdks/java/extensions/avro/src/test/java/org/apache/beam/sdk/extensions/avro/coders/AvroCoderTest.java
index a76fb59c7d10..fd13e42c0157 100644
--- a/sdks/java/extensions/avro/src/test/java/org/apache/beam/sdk/extensions/avro/coders/AvroCoderTest.java
+++ b/sdks/java/extensions/avro/src/test/java/org/apache/beam/sdk/extensions/avro/coders/AvroCoderTest.java
@@ -22,6 +22,8 @@
 import static org.hamcrest.Matchers.equalTo;
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotSame;
+import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
@@ -51,6 +53,7 @@
 import org.apache.avro.generic.GenericData;
 import org.apache.avro.generic.GenericRecord;
 import org.apache.avro.io.DatumReader;
+import org.apache.avro.io.DatumWriter;
 import org.apache.avro.reflect.AvroName;
 import org.apache.avro.reflect.AvroSchema;
 import org.apache.avro.reflect.ReflectData;
@@ -416,6 +419,95 @@ public void testReflectRecordEncoding() throws Exception {
     CoderProperties.coderDecodeEncodeEqual(coderWithSchema, AVRO_SPECIFIC_RECORD);
   }
 
+  @Test
+  public void testCoderCached() {
+    final Schema avroSchema = AVRO_SPECIFIC_RECORD.getSchema();
+    final TypeDescriptor<TestAvro> typeDescriptor = TypeDescriptor.of(TestAvro.class);
+    final Class<TestAvro> testAvroClass = TestAvro.class;
+
+    AvroDatumFactory<TestAvro> reflectFactory1 =
+        new AvroDatumFactory.ReflectDatumFactory<>(testAvroClass);
+    AvroDatumFactory<TestAvro> reflectFactory2 =
+        new AvroDatumFactory.ReflectDatumFactory<>(testAvroClass);
+    AvroDatumFactory<TestAvro> specificDatumFactory =
+        new AvroDatumFactory.SpecificDatumFactory<>(testAvroClass);
+    AvroDatumFactory<TestAvro> customDatumFactory =
+        new AvroDatumFactory<TestAvro>(TestAvro.class) {
+          @Override
+          public DatumWriter<TestAvro> apply(Schema writer) {
+            return null;
+          }
+
+          @Override
+          public DatumReader<TestAvro> apply(Schema writer, Schema reader) {
+            return null;
+          }
+        };
+
+    // Test that the AvroCoder caches the coder for the same class/type descriptor/schema/datum
+    // factory/useReflectApi
+    assertSame(AvroCoder.of(testAvroClass), AvroCoder.of(testAvroClass));
+    assertSame(AvroCoder.of(testAvroClass, false), AvroCoder.of(testAvroClass, false));
+    assertNotSame(AvroCoder.of(testAvroClass, false), AvroCoder.of(testAvroClass, true));
+    assertSame(AvroCoder.of(typeDescriptor), AvroCoder.of(typeDescriptor));
+    assertSame(AvroCoder.of(typeDescriptor, false), AvroCoder.of(typeDescriptor, false));
+    assertSame(AvroCoder.of(testAvroClass, avroSchema), AvroCoder.of(testAvroClass, avroSchema));
+    assertSame(
+        AvroCoder.of(testAvroClass, avroSchema, false),
+        AvroCoder.of(testAvroClass, avroSchema, false));
+
+    // Test that the AvroCoder caches the coder when providing a datum factory
+    assertSame(
+        AvroCoder.of(reflectFactory1, avroSchema), AvroCoder.of(reflectFactory1, avroSchema));
+    assertSame(
+        AvroCoder.of(reflectFactory1, avroSchema), AvroCoder.of(reflectFactory2, avroSchema));
+
+    // ensure that when providing datum factories of different types,
+    // the previously cached coder is not used
+    assertNotSame(
+        AvroCoder.of(reflectFactory1, avroSchema), AvroCoder.of(specificDatumFactory, avroSchema));
+
+    // for custom datum factories (neither specific nor reflect) we don't use the cache
+    assertNotSame(
+        AvroCoder.of(customDatumFactory, avroSchema), AvroCoder.of(customDatumFactory, avroSchema));
+
+    // Test that the AvroCoder caches the coder when using a specific datum factory
+    assertSame(AvroCoder.specific(testAvroClass), AvroCoder.specific(testAvroClass));
+    assertSame(AvroCoder.specific(typeDescriptor), AvroCoder.specific(typeDescriptor));
+    assertSame(
+        AvroCoder.specific(testAvroClass, avroSchema),
+        AvroCoder.specific(testAvroClass, avroSchema));
+
+    // Test that the AvroCoder caches the coder when using a reflect datum factory
+    assertSame(AvroCoder.reflect(testAvroClass), AvroCoder.reflect(testAvroClass));
+    assertSame(AvroCoder.reflect(typeDescriptor), AvroCoder.reflect(typeDescriptor));
+    assertSame(
+        AvroCoder.reflect(testAvroClass, avroSchema), AvroCoder.reflect(testAvroClass, avroSchema));
+
+    // Test that the AvroCoder does not cache the coder when using different datum factories
+    assertNotSame(AvroCoder.specific(testAvroClass), AvroCoder.reflect(testAvroClass));
+  }
+
+  @Test
+  public void testAvroGenericCoderSame() {
+    Schema schema = AVRO_SPECIFIC_RECORD.getSchema();
+    assertSame(AvroCoder.of(schema), AvroCoder.of(schema));
+    assertSame(AvroCoder.generic(schema), AvroCoder.generic(schema));
+  }
+
+  @Test
+  public void testAvroCoderCacheUsedWhenCoderDeserialized() {
+    AvroCoder<TestAvro> coder = AvroCoder.of(TestAvro.class);
+    AvroCoder<TestAvro> deserializedCoder1 = SerializableUtils.clone(coder);
+    AvroCoder<TestAvro> deserializedCoder2 = SerializableUtils.clone(coder);
+
+    assertSame(coder.getDatumReader(), deserializedCoder1.getDatumReader());
+    assertSame(coder.getDatumWriter(), deserializedCoder1.getDatumWriter());
+
+    assertSame(coder.getDatumReader(), deserializedCoder2.getDatumReader());
+    assertSame(coder.getDatumWriter(), deserializedCoder2.getDatumWriter());
+  }
+
   @Test
   public void testGenericRecordEncoding() throws Exception {
     String schemaString =


Reviews:
user: login: scwhittle
body: Thanks!
state: CHANGES_REQUESTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T08:54:44Z
user: login: wollowizard
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T12:10:09Z
user: login: wollowizard
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T12:11:01Z
user: login: wollowizard
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T12:11:11Z
user: login: scwhittle
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-08T14:38:53Z
user: login: wollowizard
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-09T08:53:38Z
user: login: scwhittle
state: CHANGES_REQUESTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-09T09:24:17Z
user: login: wollowizard
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-09T11:27:24Z
user: login: scwhittle
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-09T13:52:12Z
user: login: scwhittle
body: Thanks for identifying this problem and fixing it!
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-12T08:14:07Z

Review comments:
diff_hunk: @@ -242,7 +252,18 @@ public static <T> AvroCoder<T> of(Class<T> type, Schema schema) {
    * @param <T> the element type
    */
   public static <T> AvroCoder<T> of(AvroDatumFactory<T> datumFactory, Schema schema) {
-    return new AvroCoder<>(datumFactory, schema);
+    Class<T> type = datumFactory.getType();
+    return fromCacheOrCreate(type, schema, () -> new AvroCoder<>(datumFactory, schema));
+  }
+
+  private static <T> AvroCoder<T> fromCacheOrCreate(
+      Class<T> type, Schema schema, Callable<AvroCoder<T>> avroCoderCreator) {
path: sdks/java/extensions/avro/src/main/java/org/apache/beam/sdk/extensions/avro/coders/AvroCoder.java
user: login: scwhittle
body: it might be safer to have this take a type ENUM as well (specific, reflect). Otherwise AvroCoder.specific() factory could return something previously cached via AvroCoder.reflect() and vice versa.

That might not be expected usage but seems like we might as well guard against that since it woudl be pretty confusing.
created_at: 2025-05-07T08:19:54Z
author_association: CONTRIBUTOR
diff_hunk: @@ -840,4 +843,38 @@ public boolean equals(@Nullable Object other) {
   public int hashCode() {
     return Objects.hash(getClass(), typeDescriptor, datumFactory, schemaSupplier.get());
   }
+
+  private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundException {
path: sdks/java/extensions/avro/src/main/java/org/apache/beam/sdk/extensions/avro/coders/AvroCoder.java
user: login: scwhittle
body: Actually I think that this could still lead to duplicate reader/writers if we deserialize the same logical coder multiple times.

I think that could be fixed by using readResolve which allows modifying the deserialized object (just learning about this myself). https://docs.oracle.com/javase/8/docs/platform/serialization/spec/input.html#a5903

One idea that seems clean is to change to serialize the cache key for AvroCoder using writeReplace and then use readResolve of that key to lookup/create the coder in the cache based upon that key.  There is an example in this file of that with SerializableSchemaString.

created_at: 2025-05-07T08:34:05Z
author_association: CONTRIBUTOR
diff_hunk: @@ -17,17 +17,27 @@
  */
 package org.apache.beam.sdk.extensions.avro.coders;
 
+import java.util.concurrent.ExecutionException;
 import org.apache.avro.Schema;
 import org.apache.avro.generic.GenericRecord;
 import org.apache.beam.sdk.extensions.avro.io.AvroDatumFactory;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.Cache;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheBuilder;
 
 /** AvroCoder specialisation for GenericRecord, needed for cross-language transforms. */
 public class AvroGenericCoder extends AvroCoder<GenericRecord> {
+  private static final Cache<Schema, AvroGenericCoder> AVRO_GENERIC_CODER_CACHE =
+      CacheBuilder.newBuilder().build();
path: sdks/java/extensions/avro/src/main/java/org/apache/beam/sdk/extensions/avro/coders/AvroGenericCoder.java
user: login: scwhittle
body: use the weakValues option?  I'm not sure if there are cases where schemas are created dynamically and this cache could grow unbounded.
created_at: 2025-05-07T08:35:23Z
author_association: CONTRIBUTOR
diff_hunk: @@ -416,6 +417,40 @@ public void testReflectRecordEncoding() throws Exception {
     CoderProperties.coderDecodeEncodeEqual(coderWithSchema, AVRO_SPECIFIC_RECORD);
   }
 
+  @Test
+  public void testCoderCached() {
+    Schema schema = AVRO_SPECIFIC_RECORD.getSchema();
+    TypeDescriptor<TestAvro> typeDescriptor = TypeDescriptor.of(TestAvro.class);
+    Class<TestAvro> clazz = TestAvro.class;
+    boolean useReflectApi = false;
+    AvroDatumFactory<TestAvro> datumFactory = new AvroDatumFactory.ReflectDatumFactory<>(clazz);
+
+    assertSame(AvroCoder.of(clazz), AvroCoder.of(clazz));
+    assertSame(AvroCoder.of(clazz, useReflectApi), AvroCoder.of(clazz, useReflectApi));
+    assertSame(AvroCoder.of(typeDescriptor), AvroCoder.of(typeDescriptor));
+    assertSame(
+        AvroCoder.of(typeDescriptor, useReflectApi), AvroCoder.of(typeDescriptor, useReflectApi));
+    assertSame(AvroCoder.of(clazz, schema), AvroCoder.of(clazz, schema));
+    assertSame(
+        AvroCoder.of(clazz, schema, useReflectApi), AvroCoder.of(clazz, schema, useReflectApi));
+    assertSame(AvroCoder.of(datumFactory, schema), AvroCoder.of(datumFactory, schema));
+
+    assertSame(AvroCoder.specific(clazz), AvroCoder.specific(clazz));
path: sdks/java/extensions/avro/src/test/java/org/apache/beam/sdk/extensions/avro/coders/AvroCoderTest.java
user: login: scwhittle
body: assert that AvroCoder.specific(clazz) is different than AvroCoder.reflect(clazz)
created_at: 2025-05-07T08:36:14Z
author_association: CONTRIBUTOR
diff_hunk: @@ -840,4 +843,38 @@ public boolean equals(@Nullable Object other) {
   public int hashCode() {
     return Objects.hash(getClass(), typeDescriptor, datumFactory, schemaSupplier.get());
   }
+
+  private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundException {
path: sdks/java/extensions/avro/src/main/java/org/apache/beam/sdk/extensions/avro/coders/AvroCoder.java
user: login: wollowizard
body: @scwhittle you were right when deserializing the same coder, we would get different readers/writers.
I've fixed that in a way similar to the one you described, but maybe a bit simpler. I've added a test for that too.

As curiosity, are coders serialized and deserialized often? my assumption was that they this happens when a new vm is started for example, or when checkpointing. or does it happen even more often so to have a performance impact?
created_at: 2025-05-07T12:10:09Z
author_association: CONTRIBUTOR
diff_hunk: @@ -416,6 +417,40 @@ public void testReflectRecordEncoding() throws Exception {
     CoderProperties.coderDecodeEncodeEqual(coderWithSchema, AVRO_SPECIFIC_RECORD);
   }
 
+  @Test
+  public void testCoderCached() {
+    Schema schema = AVRO_SPECIFIC_RECORD.getSchema();
+    TypeDescriptor<TestAvro> typeDescriptor = TypeDescriptor.of(TestAvro.class);
+    Class<TestAvro> clazz = TestAvro.class;
+    boolean useReflectApi = false;
+    AvroDatumFactory<TestAvro> datumFactory = new AvroDatumFactory.ReflectDatumFactory<>(clazz);
+
+    assertSame(AvroCoder.of(clazz), AvroCoder.of(clazz));
+    assertSame(AvroCoder.of(clazz, useReflectApi), AvroCoder.of(clazz, useReflectApi));
+    assertSame(AvroCoder.of(typeDescriptor), AvroCoder.of(typeDescriptor));
+    assertSame(
+        AvroCoder.of(typeDescriptor, useReflectApi), AvroCoder.of(typeDescriptor, useReflectApi));
+    assertSame(AvroCoder.of(clazz, schema), AvroCoder.of(clazz, schema));
+    assertSame(
+        AvroCoder.of(clazz, schema, useReflectApi), AvroCoder.of(clazz, schema, useReflectApi));
+    assertSame(AvroCoder.of(datumFactory, schema), AvroCoder.of(datumFactory, schema));
+
+    assertSame(AvroCoder.specific(clazz), AvroCoder.specific(clazz));
path: sdks/java/extensions/avro/src/test/java/org/apache/beam/sdk/extensions/avro/coders/AvroCoderTest.java
user: login: wollowizard
body: done see last check in this test
created_at: 2025-05-07T12:11:01Z
author_association: CONTRIBUTOR
diff_hunk: @@ -17,17 +17,27 @@
  */
 package org.apache.beam.sdk.extensions.avro.coders;
 
+import java.util.concurrent.ExecutionException;
 import org.apache.avro.Schema;
 import org.apache.avro.generic.GenericRecord;
 import org.apache.beam.sdk.extensions.avro.io.AvroDatumFactory;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.Cache;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheBuilder;
 
 /** AvroCoder specialisation for GenericRecord, needed for cross-language transforms. */
 public class AvroGenericCoder extends AvroCoder<GenericRecord> {
+  private static final Cache<Schema, AvroGenericCoder> AVRO_GENERIC_CODER_CACHE =
+      CacheBuilder.newBuilder().build();
path: sdks/java/extensions/avro/src/main/java/org/apache/beam/sdk/extensions/avro/coders/AvroGenericCoder.java
user: login: wollowizard
body: done
created_at: 2025-05-07T12:11:11Z
author_association: CONTRIBUTOR
diff_hunk: @@ -840,4 +843,38 @@ public boolean equals(@Nullable Object other) {
   public int hashCode() {
     return Objects.hash(getClass(), typeDescriptor, datumFactory, schemaSupplier.get());
   }
+
+  private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundException {
path: sdks/java/extensions/avro/src/main/java/org/apache/beam/sdk/extensions/avro/coders/AvroCoder.java
user: login: scwhittle
body: I know that for DoFn objects we deserialize them multiple times and cache them since they are single-threaded and we run in multiple threads in parallel.  However I'm not sure if Coders are deserialized similarly.  Since they are thread-safe it seems they could be deserialized once and shared.
@kennknowles Do you know how coders are represented and deserialized off the top of your head? Should we worry about multiple copies  not sharing state or is that unexpected? Are coders represented with some set of coders for the pipeline or could the same coder perhaps be serialized multiple times in the graph?

One benefit of the writeReplace and readResolve idea is that the same AvroCoder instance would be used from multiple deserializations so the thread-local encode/decode buffers it contains would be shared, not just the reader and writer. But we could do that as a follow up.

Were you able to test this to verify it was working as expected?
created_at: 2025-05-08T14:38:53Z
author_association: CONTRIBUTOR
diff_hunk: @@ -840,4 +843,38 @@ public boolean equals(@Nullable Object other) {
   public int hashCode() {
     return Objects.hash(getClass(), typeDescriptor, datumFactory, schemaSupplier.get());
   }
+
+  private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundException {
path: sdks/java/extensions/avro/src/main/java/org/apache/beam/sdk/extensions/avro/coders/AvroCoder.java
user: login: wollowizard
body: @scwhittle yes I could test it with realistic data, and it works as expected.

**AvroDatumFactory**
before: 32%
after: 0.036%

**AvroCoder**
before: 69%
after: 11%


As we thought, this improves not only user code, but also sdk performance.
Example taken from another pipeline:
**org.apache.beam.runners.dataflow.worker.StreamingModeExecutionContext.flushState()**
before: 43%
after: 10%


**Before**

<img width="600" alt="image" src="https://github.com/user-attachments/assets/b121e360-79c9-4fbd-b9d9-079ed9ad50c3" />
<img width="600" alt="image" src="https://github.com/user-attachments/assets/3a0e3651-c5f3-4d0b-a2ef-87ec61733bfe" />
<img width="600" alt="image" src="https://github.com/user-attachments/assets/6173ea0e-a9b6-4a23-9653-d62df14ef61e" />





**After**

<img width="600" alt="image" src="https://github.com/user-attachments/assets/a33c71f6-af23-481e-a721-5fb272161e67" />
<img width="600" alt="image" src="https://github.com/user-attachments/assets/31ee28a8-637c-4bd6-a95e-7eca0b19040b" />
<img width="600" alt="image" src="https://github.com/user-attachments/assets/ae1516dc-d23a-4857-ae09-e2c28e4977a0" />



created_at: 2025-05-09T08:53:37Z
author_association: CONTRIBUTOR
diff_hunk: @@ -17,17 +17,27 @@
  */
 package org.apache.beam.sdk.extensions.avro.coders;
 
+import java.util.concurrent.ExecutionException;
 import org.apache.avro.Schema;
 import org.apache.avro.generic.GenericRecord;
 import org.apache.beam.sdk.extensions.avro.io.AvroDatumFactory;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.Cache;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheBuilder;
 
 /** AvroCoder specialisation for GenericRecord, needed for cross-language transforms. */
 public class AvroGenericCoder extends AvroCoder<GenericRecord> {
+  private static final Cache<Schema, AvroGenericCoder> AVRO_GENERIC_CODER_CACHE =
+      CacheBuilder.newBuilder().build();
path: sdks/java/extensions/avro/src/main/java/org/apache/beam/sdk/extensions/avro/coders/AvroGenericCoder.java
user: login: scwhittle
body: actually this cache could still be weakValues
created_at: 2025-05-09T09:01:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -354,8 +392,10 @@ public Schema get() {
   // an inner coder.
   private final EmptyOnDeserializationThreadLocal<BinaryDecoder> decoder;
   private final EmptyOnDeserializationThreadLocal<BinaryEncoder> encoder;
-  private final EmptyOnDeserializationThreadLocal<DatumWriter<T>> writer;
path: sdks/java/extensions/avro/src/main/java/org/apache/beam/sdk/extensions/avro/coders/AvroCoder.java
user: login: scwhittle
body: Sorry I didn't think of this earlier, but I realized that this is going to change the java serialization of this coder which will lead to update incompatability.

Options:
- keep the existing fields and the new transient fields and just don't use existing fields.
- remove the new fields and change existing fields to do the cache lookup in the initialvalue method.  I think you could remove readObject then. We are using unnecessary per-thread caching still then.

I think I'd lean towards the first option of keeping the existing fields but not using them.
created_at: 2025-05-09T09:23:07Z
author_association: CONTRIBUTOR
diff_hunk: @@ -840,4 +843,38 @@ public boolean equals(@Nullable Object other) {
   public int hashCode() {
     return Objects.hash(getClass(), typeDescriptor, datumFactory, schemaSupplier.get());
   }
+
+  private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundException {
path: sdks/java/extensions/avro/src/main/java/org/apache/beam/sdk/extensions/avro/coders/AvroCoder.java
user: login: scwhittle
body: This is great, thanks!
created_at: 2025-05-09T09:23:58Z
author_association: CONTRIBUTOR
diff_hunk: @@ -354,8 +392,10 @@ public Schema get() {
   // an inner coder.
   private final EmptyOnDeserializationThreadLocal<BinaryDecoder> decoder;
   private final EmptyOnDeserializationThreadLocal<BinaryEncoder> encoder;
-  private final EmptyOnDeserializationThreadLocal<DatumWriter<T>> writer;
path: sdks/java/extensions/avro/src/main/java/org/apache/beam/sdk/extensions/avro/coders/AvroCoder.java
user: login: wollowizard
body: @scwhittle i see, i guess this is to support [pipeline upgrades](https://cloud.google.com/dataflow/docs/guides/updating-a-pipeline)?
anyway I restored the old fields 
created_at: 2025-05-09T11:27:24Z
author_association: CONTRIBUTOR
diff_hunk: @@ -416,6 +419,88 @@ public void testReflectRecordEncoding() throws Exception {
     CoderProperties.coderDecodeEncodeEqual(coderWithSchema, AVRO_SPECIFIC_RECORD);
   }
 
+  @Test
+  public void testCoderCached() {
+    Schema schema = AVRO_SPECIFIC_RECORD.getSchema();
+    TypeDescriptor<TestAvro> typeDescriptor = TypeDescriptor.of(TestAvro.class);
+    Class<TestAvro> clazz = TestAvro.class;
+    boolean useReflectApi = false;
path: sdks/java/extensions/avro/src/test/java/org/apache/beam/sdk/extensions/avro/coders/AvroCoderTest.java
user: login: scwhittle
body: nit: how about 
final boolean USE_REFLECT_API = true;

I was a little confused about this at first thinking you meant to test with both true/false. But it's also odd to set this to false if the parameter is useReflectApi as it is then inverted.
created_at: 2025-05-09T13:52:12Z
author_association: CONTRIBUTOR
diff_hunk: @@ -354,8 +392,10 @@ public Schema get() {
   // an inner coder.
   private final EmptyOnDeserializationThreadLocal<BinaryDecoder> decoder;
   private final EmptyOnDeserializationThreadLocal<BinaryEncoder> encoder;
-  private final EmptyOnDeserializationThreadLocal<DatumWriter<T>> writer;
path: sdks/java/extensions/avro/src/main/java/org/apache/beam/sdk/extensions/avro/coders/AvroCoder.java
user: login: scwhittle
body: Yes, this is for compatability across updates in which the pipeline has been reserialized and may be before and after this change.

There is an updateCompatabilityVersion option that can be used to guard uncompatible changes but seems simpler in this case to just leave the old unused fields.
created_at: 2025-05-12T08:12:25Z
author_association: CONTRIBUTOR

PR comments:
user: login: wollowizard
created_at: 2025-05-06T21:20:17Z
body: assign set of reviewers
author_association: CONTRIBUTOR
user: login: wollowizard
created_at: 2025-05-06T21:22:34Z
body: @scwhittle this is a follow up to our discussion in https://github.com/apache/beam/issues/34749 and https://github.com/apache/beam/pull/34750. I would like to run a test with real data in the dataflow runner, is there a way to get a snapshot build for this PR? 
author_association: CONTRIBUTOR
user: login: scwhittle
created_at: 2025-05-07T08:54:30Z
body: > @scwhittle this is a follow up to our discussion in #34749 and #34750. I would like to run a test with real data in the dataflow runner, is there a way to get a snapshot build for this PR?

I think that you should be able to run a pipeline with these changes by publishing the jars to your local maven and then configuring your pipeline to use them.

# In the repo with this PR
./gradlew sdks:java:io:google-cloud-platform:compileJava \
&& ./gradlew -Ppublishing -p sdks/java/io/google-cloud-platform publishToMavenLocal
# find local maven path
mvn help:evaluate -Dexpression=settings.localRepository

# In the project with your pipeline, modify pom.xml to use local repo
<repositories>
    <repository>
      <id>example-repo</id>
      <name>Example Repository</name>
      <url>file:///complete/path/to/.m2/repository</url>
    </repository>
    ...
</repositories>
# Also modify the BEAM version in the pom.xml to match the version published above. ie (<beam.version>2.64.0</beam.version>) to the SNAPSHOT version (<beam.version>2.65.0-SNAPSHOT</beam.version>)

Then if you launch the dataflow pipeline as normal from your pipeline it should pull in the locally published sdk changes. Let me know if that works or if you have problems with that
author_association: CONTRIBUTOR
user: login: github-actions[bot]
created_at: 2025-05-07T13:16:24Z
body: Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment `assign set of reviewers`
author_association: CONTRIBUTOR
user: login: scwhittle
created_at: 2025-05-09T13:58:56Z
body: The test failures look unrelated I will retrigger just to be safe though
author_association: CONTRIBUTOR
user: login: scwhittle
created_at: 2025-05-09T13:59:06Z
body: Run Java PreCommit
author_association: CONTRIBUTOR
user: login: github-actions[bot]
created_at: 2025-05-10T07:06:35Z
body: Assigning reviewers:

R: @kennknowles for label java.


Note: If you would like to opt out of this review, comment `assign to next reviewer`.

Available commands:
- `stop reviewer notifications` - opt out of the automated review tooling
- `remind me after tests pass` - tag the comment author after tests pass
- `waiting on author` - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).
author_association: CONTRIBUTOR