Make SerializableConfiguration cacheable

ddebowczyk92 · Sep 26, 2023 · b233609 · b233609
1 parent d52b077
commit b233609
Showing 1 changed file with 15 additions and 1 deletion.
diff --git a/.../hadoop-common/src/main/java/org/apache/beam/sdk/io/hadoop/SerializableConfiguration.java b/.../hadoop-common/src/main/java/org/apache/beam/sdk/io/hadoop/SerializableConfiguration.java
@@ -17,6 +17,8 @@
  */
 package org.apache.beam.sdk.io.hadoop;
 
+import java.io.ByteArrayOutputStream;
+import java.io.DataOutputStream;
 import java.io.Externalizable;
 import java.io.IOException;
 import java.io.ObjectInput;
@@ -39,6 +41,8 @@ public class SerializableConfiguration implements Externalizable {
 
   private transient Configuration conf;
 
+  private transient byte[] serializationCache;
+
   public SerializableConfiguration() {}
 
   public SerializableConfiguration(Configuration conf) {
@@ -49,13 +53,23 @@ public SerializableConfiguration(Configuration conf) {
   }
 
   public Configuration get() {
+    if(serializationCache != null) {
+      serializationCache = null;
+    }
     return conf;
   }
 
   @Override
   public void writeExternal(ObjectOutput out) throws IOException {
+    if (serializationCache == null) {
+      ByteArrayOutputStream baos = new ByteArrayOutputStream(512);
+      try (DataOutputStream dos = new DataOutputStream(baos)) {
+        conf.write(dos);
+        serializationCache = baos.toByteArray();
+      }
+    }
     out.writeUTF(conf.getClass().getCanonicalName());
-    conf.write(out);
+    out.write(serializationCache);
   }
 
   @Override