InMemoryStateInternals.copy clones the values using the coder

Author: jkff

runners/core-java/src/main/java/org/apache/beam/runners/core/InMemoryStateInternals.java

@@ -34,6 +34,7 @@
 import org.apache.beam.sdk.annotations.Experimental;
 import org.apache.beam.sdk.annotations.Experimental.Kind;
 import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderException;
 import org.apache.beam.sdk.state.BagState;
 import org.apache.beam.sdk.state.CombiningState;
 import org.apache.beam.sdk.state.MapState;
@@ -49,6 +50,7 @@
 import org.apache.beam.sdk.transforms.CombineWithContext.CombineFnWithContext;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.transforms.windowing.TimestampCombiner;
+import org.apache.beam.sdk.util.CoderUtils;
 import org.apache.beam.sdk.util.CombineFnUtil;
 import org.joda.time.Instant;
 
@@ -126,25 +128,25 @@ public InMemoryStateBinder(StateContext<?> c) {
     @Override
     public <T> ValueState<T> bindValue(
         StateTag<ValueState<T>> address, Coder<T> coder) {
-      return new InMemoryValue<>();
+      return new InMemoryValue<>(coder);
     }
 
     @Override
     public <T> BagState<T> bindBag(
         final StateTag<BagState<T>> address, Coder<T> elemCoder) {
-      return new InMemoryBag<>();
+      return new InMemoryBag<>(elemCoder);
     }
 
     @Override
     public <T> SetState<T> bindSet(StateTag<SetState<T>> spec, Coder<T> elemCoder) {
-      return new InMemorySet<>();
+      return new InMemorySet<>(elemCoder);
     }
 
     @Override
     public <KeyT, ValueT> MapState<KeyT, ValueT> bindMap(
         StateTag<MapState<KeyT, ValueT>> spec,
         Coder<KeyT> mapKeyCoder, Coder<ValueT> mapValueCoder) {
-      return new InMemoryMap<>();
+      return new InMemoryMap<>(mapKeyCoder, mapValueCoder);
     }
 
     @Override
@@ -153,7 +155,7 @@ public <KeyT, ValueT> MapState<KeyT, ValueT> bindMap(
             StateTag<CombiningState<InputT, AccumT, OutputT>> address,
             Coder<AccumT> accumCoder,
             final CombineFn<InputT, AccumT, OutputT> combineFn) {
-      return new InMemoryCombiningState<>(combineFn);
+      return new InMemoryCombiningState<>(combineFn, accumCoder);
     }
 
     @Override
@@ -178,9 +180,15 @@ public WatermarkHoldState bindWatermark(
    */
   public static final class InMemoryValue<T>
       implements ValueState<T>, InMemoryState<InMemoryValue<T>> {
+    private final Coder<T> coder;
+
     private boolean isCleared = true;
     private @Nullable T value = null;
 
+    public InMemoryValue(Coder<T> coder) {
+      this.coder = coder;
+    }
+
     @Override
     public void clear() {
       // Even though we're clearing we can't remove this from the in-memory state map, since
@@ -207,10 +215,10 @@ public void write(T input) {
 
     @Override
     public InMemoryValue<T> copy() {
-      InMemoryValue<T> that = new InMemoryValue<>();
+      InMemoryValue<T> that = new InMemoryValue<>(coder);
       if (!this.isCleared) {
         that.isCleared = this.isCleared;
-        that.value = this.value;
+        that.value = unsafeClone(coder, this.value);
       }
       return that;
     }
@@ -305,14 +313,16 @@ public InMemoryWatermarkHold<W> copy() {
   public static final class InMemoryCombiningState<InputT, AccumT, OutputT>
       implements CombiningState<InputT, AccumT, OutputT>,
           InMemoryState<InMemoryCombiningState<InputT, AccumT, OutputT>> {
-    private boolean isCleared = true;
     private final CombineFn<InputT, AccumT, OutputT> combineFn;
+    private final Coder<AccumT> accumCoder;
+    private boolean isCleared = true;
     private AccumT accum;
 
     public InMemoryCombiningState(
-        CombineFn<InputT, AccumT, OutputT> combineFn) {
+        CombineFn<InputT, AccumT, OutputT> combineFn, Coder<AccumT> accumCoder) {
       this.combineFn = combineFn;
       accum = combineFn.createAccumulator();
+      this.accumCoder = accumCoder;
     }
 
     @Override
@@ -378,7 +388,7 @@ public boolean isCleared() {
     @Override
     public InMemoryCombiningState<InputT, AccumT, OutputT> copy() {
       InMemoryCombiningState<InputT, AccumT, OutputT> that =
-          new InMemoryCombiningState<>(combineFn);
+          new InMemoryCombiningState<>(combineFn, accumCoder);
       if (!this.isCleared) {
         that.isCleared = this.isCleared;
         that.addAccum(accum);
@@ -391,8 +401,13 @@ public InMemoryCombiningState<InputT, AccumT, OutputT> copy() {
    * An {@link InMemoryState} implementation of {@link BagState}.
    */
   public static final class InMemoryBag<T> implements BagState<T>, InMemoryState<InMemoryBag<T>> {
+    private final Coder<T> elemCoder;
     private List<T> contents = new ArrayList<>();
 
+    public InMemoryBag(Coder<T> elemCoder) {
+      this.elemCoder = elemCoder;
+    }
+
     @Override
     public void clear() {
       // Even though we're clearing we can't remove this from the in-memory state map, since
@@ -442,8 +457,10 @@ public Boolean read() {
 
     @Override
     public InMemoryBag<T> copy() {
-      InMemoryBag<T> that = new InMemoryBag<>();
-      that.contents.addAll(this.contents);
+      InMemoryBag<T> that = new InMemoryBag<>(elemCoder);
+      for (T elem : this.contents) {
+        that.contents.add(unsafeClone(elemCoder, elem));
+      }
       return that;
     }
   }
@@ -452,8 +469,13 @@ public InMemoryBag<T> copy() {
    * An {@link InMemoryState} implementation of {@link SetState}.
    */
   public static final class InMemorySet<T> implements SetState<T>, InMemoryState<InMemorySet<T>> {
+    private final Coder<T> elemCoder;
     private Set<T> contents = new HashSet<>();
 
+    public InMemorySet(Coder<T> elemCoder) {
+      this.elemCoder = elemCoder;
+    }
+
     @Override
     public void clear() {
       contents = new HashSet<>();
@@ -513,8 +535,10 @@ public Boolean read() {
 
     @Override
     public InMemorySet<T> copy() {
-      InMemorySet<T> that = new InMemorySet<>();
-      that.contents.addAll(this.contents);
+      InMemorySet<T> that = new InMemorySet<>(elemCoder);
+      for (T elem : this.contents) {
+        that.contents.add(unsafeClone(elemCoder, elem));
+      }
       return that;
     }
   }
@@ -524,8 +548,16 @@ public InMemorySet<T> copy() {
    */
   public static final class InMemoryMap<K, V> implements
       MapState<K, V>, InMemoryState<InMemoryMap<K, V>> {
+    private final Coder<K> keyCoder;
+    private final Coder<V> valueCoder;
+
     private Map<K, V> contents = new HashMap<>();
 
+    public InMemoryMap(Coder<K> keyCoder, Coder<V> valueCoder) {
+      this.keyCoder = keyCoder;
+      this.valueCoder = valueCoder;
+    }
+
     @Override
     public void clear() {
       contents = new HashMap<>();
@@ -600,9 +632,21 @@ public boolean isCleared() {
 
     @Override
     public InMemoryMap<K, V> copy() {
-      InMemoryMap<K, V> that = new InMemoryMap<>();
+      InMemoryMap<K, V> that = new InMemoryMap<>(keyCoder, valueCoder);
+      for (Map.Entry<K, V> entry : this.contents.entrySet()) {
+        that.contents.put(
+            unsafeClone(keyCoder, entry.getKey()), unsafeClone(valueCoder, entry.getValue()));
+      }
       that.contents.putAll(this.contents);
       return that;
     }
   }
+
+  private static <T> T unsafeClone(Coder<T> coder, T value) {
+    try {
+      return CoderUtils.clone(coder, value);
+    } catch (CoderException e) {
+      throw new RuntimeException(e);
+    }
+  }
 }

runners/direct-java/src/main/java/org/apache/beam/runners/direct/CopyOnAccessInMemoryStateInternals.java

@@ -300,7 +300,7 @@ public <T> ValueState<T> bindValue(
                       underlying.get().get(namespace, address, c);
               return existingState.copy();
             } else {
-              return new InMemoryValue<>();
+              return new InMemoryValue<>(coder);
             }
           }
 
@@ -317,7 +317,7 @@ CombiningState<InputT, AccumT, OutputT> bindCombiningValue(
                       underlying.get().get(namespace, address, c);
               return existingState.copy();
             } else {
-              return new InMemoryCombiningState<>(combineFn);
+              return new InMemoryCombiningState<>(combineFn, accumCoder);
             }
           }
 
@@ -331,7 +331,7 @@ public <T> BagState<T> bindBag(
                       underlying.get().get(namespace, address, c);
               return existingState.copy();
             } else {
-              return new InMemoryBag<>();
+              return new InMemoryBag<>(elemCoder);
             }
           }
 
@@ -345,7 +345,7 @@ public <T> SetState<T> bindSet(
                       underlying.get().get(namespace, address, c);
               return existingState.copy();
             } else {
-              return new InMemorySet<>();
+              return new InMemorySet<>(elemCoder);
             }
           }
 
@@ -361,7 +361,7 @@ public <KeyT, ValueT> MapState<KeyT, ValueT> bindMap(
                       underlying.get().get(namespace, address, c);
               return existingState.copy();
             } else {
-              return new InMemoryMap<>();
+              return new InMemoryMap<>(mapKeyCoder, mapValueCoder);
             }
           }