beam-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From aviem...@apache.org
Subject [1/2] beam git commit: [BEAM-1337] Infer state coders
Date Sat, 01 Apr 2017 07:28:35 GMT
Repository: beam
Updated Branches:
  refs/heads/master 03dce6dcc -> e31ca8b0d


[BEAM-1337] Infer state coders


Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/42e690e8
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/42e690e8
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/42e690e8

Branch: refs/heads/master
Commit: 42e690e84a9f05d508f2528b1444b26ce031e080
Parents: 03dce6d
Author: Aviem Zur <aviemzur@gmail.com>
Authored: Wed Mar 1 07:27:57 2017 +0200
Committer: Aviem Zur <aviemzur@gmail.com>
Committed: Sat Apr 1 10:27:14 2017 +0300

----------------------------------------------------------------------
 .../org/apache/beam/sdk/transforms/ParDo.java   |  62 ++
 .../apache/beam/sdk/util/state/StateSpec.java   |  15 +
 .../apache/beam/sdk/util/state/StateSpecs.java  | 264 ++++++++-
 .../apache/beam/sdk/transforms/ParDoTest.java   | 578 +++++++++++++++++++
 4 files changed, 902 insertions(+), 17 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/beam/blob/42e690e8/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/ParDo.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/ParDo.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/ParDo.java
index 664fbc3..3de845b 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/ParDo.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/ParDo.java
@@ -22,6 +22,8 @@ import static com.google.common.base.Preconditions.checkArgument;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import java.io.Serializable;
+import java.lang.reflect.ParameterizedType;
+import java.lang.reflect.Type;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
@@ -29,6 +31,7 @@ import java.util.Map;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.coders.CannotProvideCoderException;
 import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
 import org.apache.beam.sdk.runners.PipelineRunner;
 import org.apache.beam.sdk.transforms.display.DisplayData;
 import org.apache.beam.sdk.transforms.display.DisplayData.Builder;
@@ -41,6 +44,7 @@ import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.transforms.windowing.WindowFn;
 import org.apache.beam.sdk.util.NameUtils;
 import org.apache.beam.sdk.util.SerializableUtils;
+import org.apache.beam.sdk.util.state.StateSpec;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.sdk.values.PCollectionTuple;
 import org.apache.beam.sdk.values.PCollectionView;
@@ -434,6 +438,59 @@ public class ParDo {
     return DisplayData.item("fn", fn.getClass()).withLabel("Transform Function");
   }
 
+  private static void finishSpecifyingStateSpecs(
+      DoFn<?, ?> fn,
+      CoderRegistry coderRegistry,
+      Coder<?> inputCoder) {
+    DoFnSignature signature = DoFnSignatures.getSignature(fn.getClass());
+    Map<String, DoFnSignature.StateDeclaration> stateDeclarations = signature.stateDeclarations();
+    for (DoFnSignature.StateDeclaration stateDeclaration : stateDeclarations.values()) {
+      try {
+        StateSpec<?, ?> stateSpec = (StateSpec<?, ?>) stateDeclaration.field().get(fn);
+        stateSpec.offerCoders(codersForStateSpecTypes(stateDeclaration, coderRegistry, inputCoder));
+        stateSpec.finishSpecifying();
+      } catch (IllegalAccessException e) {
+        throw new RuntimeException(e);
+      }
+    }
+  }
+
+  /**
+   * Try to provide coders for as many of the type arguments of given
+   * {@link DoFnSignature.StateDeclaration} as possible.
+   */
+  private static <InputT> Coder[] codersForStateSpecTypes(
+      DoFnSignature.StateDeclaration stateDeclaration,
+      CoderRegistry coderRegistry,
+      Coder<InputT> inputCoder) {
+    Type stateType = stateDeclaration.stateType().getType();
+
+    if (!(stateType instanceof ParameterizedType)) {
+      // No type arguments means no coders to infer.
+      return new Coder[0];
+    }
+
+    Type[] typeArguments = ((ParameterizedType) stateType).getActualTypeArguments();
+    Coder[] coders = new Coder[typeArguments.length];
+
+    for (int i = 0; i < typeArguments.length; i++) {
+      Type typeArgument = typeArguments[i];
+      TypeDescriptor<?> typeDescriptor = TypeDescriptor.of(typeArgument);
+      try {
+        coders[i] = coderRegistry.getDefaultCoder(typeDescriptor);
+      } catch (CannotProvideCoderException e) {
+        try {
+          coders[i] = coderRegistry.getDefaultCoder(
+              typeDescriptor, inputCoder.getEncodedTypeDescriptor(), inputCoder);
+        } catch (CannotProvideCoderException ignored) {
+          // Since not all type arguments will have a registered coder we ignore this exception.
+        }
+      }
+    }
+
+    return coders;
+  }
+
   /**
    * Perform common validations of the {@link DoFn} against the input {@link PCollection}, for
    * example ensuring that the window type expected by the {@link DoFn} matches the window type of
@@ -557,6 +614,7 @@ public class ParDo {
 
     @Override
     public PCollection<OutputT> expand(PCollection<? extends InputT> input) {
+      finishSpecifyingStateSpecs(fn, input.getPipeline().getCoderRegistry(), input.getCoder());
       TupleTag<OutputT> mainOutput = new TupleTag<>();
       return input.apply(withOutputTags(mainOutput, TupleTagList.empty())).get(mainOutput);
     }
@@ -681,6 +739,10 @@ public class ParDo {
     public PCollectionTuple expand(PCollection<? extends InputT> input) {
       // SplittableDoFn should be forbidden on the runner-side.
       validateWindowType(input, fn);
+
+      // Use coder registry to determine coders for all StateSpec defined in the fn signature.
+      finishSpecifyingStateSpecs(fn, input.getPipeline().getCoderRegistry(), input.getCoder());
+
       PCollectionTuple outputs = PCollectionTuple.ofPrimitiveOutputsInternal(
           input.getPipeline(),
           TupleTagList.of(mainOutputTag).and(sideOutputTags.getAll()),

http://git-wip-us.apache.org/repos/asf/beam/blob/42e690e8/sdks/java/core/src/main/java/org/apache/beam/sdk/util/state/StateSpec.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/state/StateSpec.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/state/StateSpec.java
index 4fdeefb..6b94c40 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/state/StateSpec.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/state/StateSpec.java
@@ -20,6 +20,7 @@ package org.apache.beam.sdk.util.state;
 import java.io.Serializable;
 import org.apache.beam.sdk.annotations.Experimental;
 import org.apache.beam.sdk.annotations.Experimental.Kind;
+import org.apache.beam.sdk.coders.Coder;
 
 /**
  * A specification of a persistent state cell. This includes information necessary to encode the
@@ -36,4 +37,18 @@ public interface StateSpec<K, StateT extends State> extends Serializable {
    * Use the {@code binder} to create an instance of {@code StateT} appropriate for this address.
    */
   StateT bind(String id, StateBinder<? extends K> binder);
+
+  /**
+   * Given {code coders} are inferred from type arguments defined for this class.
+   * Coders which are already set should take precedence over offered coders.
+   * @param coders Array of coders indexed by the type arguments order. Entries might be null if
+   *               the coder could not be inferred.
+   */
+  void offerCoders(Coder[] coders);
+
+  /**
+   * Validates that this {@link StateSpec} has been specified correctly and finalizes it.
+   * Automatically invoked when the pipeline is built.
+   */
+  void finishSpecifying();
 }

http://git-wip-us.apache.org/repos/asf/beam/blob/42e690e8/sdks/java/core/src/main/java/org/apache/beam/sdk/util/state/StateSpecs.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/state/StateSpecs.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/state/StateSpecs.java
index 8912993..6a8c80b 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/state/StateSpecs.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/state/StateSpecs.java
@@ -17,7 +17,10 @@
  */
 package org.apache.beam.sdk.util.state;
 
+import static com.google.common.base.Preconditions.checkArgument;
+
 import java.util.Objects;
+import javax.annotation.Nullable;
 import org.apache.beam.sdk.annotations.Experimental;
 import org.apache.beam.sdk.annotations.Experimental.Kind;
 import org.apache.beam.sdk.coders.CannotProvideCoderException;
@@ -44,7 +47,13 @@ public class StateSpecs {
   private StateSpecs() {}
 
   /** Create a simple state spec for values of type {@code T}. */
+  public static <T> StateSpec<Object, ValueState<T>> value() {
+    return new ValueStateSpec<>(null);
+  }
+
+  /** Create a simple state spec for values of type {@code T}. */
   public static <T> StateSpec<Object, ValueState<T>> value(Coder<T> valueCoder) {
+    checkArgument(valueCoder != null, "valueCoder should not be null. Consider value() instead");
     return new ValueStateSpec<>(valueCoder);
   }
 
@@ -53,8 +62,21 @@ public class StateSpecs {
    * {@code InputT}s into a single {@code OutputT}.
    */
   public static <InputT, AccumT, OutputT>
+  StateSpec<Object, AccumulatorCombiningState<InputT, AccumT, OutputT>> combiningValue(
+      CombineFn<InputT, AccumT, OutputT> combineFn) {
+    return new CombiningValueStateSpec<InputT, AccumT, OutputT>(null, combineFn);
+  }
+
+  /**
+   * Create a state spec for values that use a {@link CombineFn} to automatically merge multiple
+   * {@code InputT}s into a single {@code OutputT}.
+   */
+  public static <InputT, AccumT, OutputT>
       StateSpec<Object, AccumulatorCombiningState<InputT, AccumT, OutputT>> combiningValue(
           Coder<AccumT> accumCoder, CombineFn<InputT, AccumT, OutputT> combineFn) {
+    checkArgument(accumCoder != null,
+        "accumCoder should not be null. "
+            + "Consider using combiningValue(CombineFn<> combineFn) instead.");
     return combiningValueInternal(accumCoder, combineFn);
   }
 
@@ -63,8 +85,21 @@ public class StateSpecs {
    * multiple {@code InputT}s into a single {@code OutputT}.
    */
   public static <K, InputT, AccumT, OutputT>
+  StateSpec<K, AccumulatorCombiningState<InputT, AccumT, OutputT>> keyedCombiningValue(
+      KeyedCombineFn<K, InputT, AccumT, OutputT> combineFn) {
+    return new KeyedCombiningValueStateSpec<K, InputT, AccumT, OutputT>(null, combineFn);
+  }
+
+  /**
+   * Create a state spec for values that use a {@link KeyedCombineFn} to automatically merge
+   * multiple {@code InputT}s into a single {@code OutputT}.
+   */
+  public static <K, InputT, AccumT, OutputT>
       StateSpec<K, AccumulatorCombiningState<InputT, AccumT, OutputT>> keyedCombiningValue(
           Coder<AccumT> accumCoder, KeyedCombineFn<K, InputT, AccumT, OutputT> combineFn) {
+    checkArgument(accumCoder != null,
+        "accumCoder should not be null. "
+            + "Consider using keyedCombiningValue(KeyedCombineFn<> combineFn) instead.");
     return keyedCombiningValueInternal(accumCoder, combineFn);
   }
 
@@ -73,10 +108,23 @@ public class StateSpecs {
    * merge multiple {@code InputT}s into a single {@code OutputT}.
    */
   public static <K, InputT, AccumT, OutputT>
+  StateSpec<K, AccumulatorCombiningState<InputT, AccumT, OutputT>>
+  keyedCombiningValueWithContext(KeyedCombineFnWithContext<K, InputT, AccumT, OutputT> combineFn) {
+    return new KeyedCombiningValueWithContextStateSpec<K, InputT, AccumT, OutputT>(null, combineFn);
+  }
+
+  /**
+   * Create a state spec for values that use a {@link KeyedCombineFnWithContext} to automatically
+   * merge multiple {@code InputT}s into a single {@code OutputT}.
+   */
+  public static <K, InputT, AccumT, OutputT>
       StateSpec<K, AccumulatorCombiningState<InputT, AccumT, OutputT>>
           keyedCombiningValueWithContext(
               Coder<AccumT> accumCoder,
               KeyedCombineFnWithContext<K, InputT, AccumT, OutputT> combineFn) {
+    checkArgument(accumCoder != null,
+        "accumCoder should not be null. Consider using "
+            + "keyedCombiningValueWithContext(KeyedCombineFnWithContext<> combineFn) instead.");
     return new KeyedCombiningValueWithContextStateSpec<K, InputT, AccumT, OutputT>(
         accumCoder, combineFn);
   }
@@ -121,8 +169,23 @@ public class StateSpecs {
    * Create a state spec that is optimized for adding values frequently, and occasionally retrieving
    * all the values that have been added.
    */
+  public static <T> StateSpec<Object, BagState<T>> bag() {
+    return bag(null);
+  }
+
+  /**
+   * Create a state spec that is optimized for adding values frequently, and occasionally retrieving
+   * all the values that have been added.
+   */
   public static <T> StateSpec<Object, BagState<T>> bag(Coder<T> elemCoder) {
-    return new BagStateSpec<T>(elemCoder);
+    return new BagStateSpec<>(elemCoder);
+  }
+
+  /**
+   * Create a state spec that supporting for {@link java.util.Set} like access patterns.
+   */
+  public static <T> StateSpec<Object, SetState<T>> set() {
+    return set(null);
   }
 
   /**
@@ -135,6 +198,13 @@ public class StateSpecs {
   /**
    * Create a state spec that supporting for {@link java.util.Map} like access patterns.
    */
+  public static <K, V> StateSpec<Object, MapState<K, V>> map() {
+    return new MapStateSpec<>(null, null);
+  }
+
+  /**
+   * Create a state spec that supporting for {@link java.util.Map} like access patterns.
+   */
   public static <K, V> StateSpec<Object, MapState<K, V>> map(Coder<K> keyCoder,
                                                              Coder<V> valueCoder) {
     return new MapStateSpec<>(keyCoder, valueCoder);
@@ -174,9 +244,10 @@ public class StateSpecs {
    */
   private static class ValueStateSpec<T> implements StateSpec<Object, ValueState<T>> {
 
-    private final Coder<T> coder;
+    @Nullable
+    private Coder<T> coder;
 
-    private ValueStateSpec(Coder<T> coder) {
+    private ValueStateSpec(@Nullable Coder<T> coder) {
       this.coder = coder;
     }
 
@@ -185,6 +256,25 @@ public class StateSpecs {
       return visitor.bindValue(id, this, coder);
     }
 
+    @SuppressWarnings("unchecked")
+    @Override
+    public void offerCoders(Coder[] coders) {
+      if (this.coder == null) {
+        if (coders[0] != null) {
+          this.coder = (Coder<T>) coders[0];
+        }
+      }
+    }
+
+    @Override public void finishSpecifying() {
+      if (coder == null) {
+        throw new IllegalStateException("Unable to infer a coder for ValueState and no Coder"
+            + " was specified. Please set a coder by either invoking"
+            + " StateSpecs.value(Coder<T> valueCoder) or by registering the coder in the"
+            + " Pipeline's CoderRegistry.");
+      }
+    }
+
     @Override
     public boolean equals(Object obj) {
       if (obj == this) {
@@ -214,15 +304,32 @@ public class StateSpecs {
       extends KeyedCombiningValueStateSpec<Object, InputT, AccumT, OutputT>
       implements StateSpec<Object, AccumulatorCombiningState<InputT, AccumT, OutputT>> {
 
-    private final Coder<AccumT> accumCoder;
+    @Nullable
+    private Coder<AccumT> accumCoder;
     private final CombineFn<InputT, AccumT, OutputT> combineFn;
 
     private CombiningValueStateSpec(
-        Coder<AccumT> accumCoder, CombineFn<InputT, AccumT, OutputT> combineFn) {
+        @Nullable Coder<AccumT> accumCoder,
+        CombineFn<InputT, AccumT, OutputT> combineFn) {
       super(accumCoder, combineFn.asKeyedFn());
       this.combineFn = combineFn;
       this.accumCoder = accumCoder;
     }
+
+    @Override
+    protected Coder<AccumT> getAccumCoder() {
+      return accumCoder;
+    }
+
+    @SuppressWarnings("unchecked")
+    @Override
+    public void offerCoders(Coder[] coders) {
+      if (this.accumCoder == null) {
+        if (coders[1] != null) {
+          this.accumCoder = (Coder<AccumT>) coders[1];
+        }
+      }
+    }
   }
 
   /**
@@ -234,11 +341,13 @@ public class StateSpecs {
   private static class KeyedCombiningValueWithContextStateSpec<K, InputT, AccumT, OutputT>
       implements StateSpec<K, AccumulatorCombiningState<InputT, AccumT, OutputT>> {
 
-    private final Coder<AccumT> accumCoder;
+    @Nullable
+    private Coder<AccumT> accumCoder;
     private final KeyedCombineFnWithContext<K, InputT, AccumT, OutputT> combineFn;
 
     protected KeyedCombiningValueWithContextStateSpec(
-        Coder<AccumT> accumCoder, KeyedCombineFnWithContext<K, InputT, AccumT, OutputT> combineFn) {
+        @Nullable Coder<AccumT> accumCoder,
+        KeyedCombineFnWithContext<K, InputT, AccumT, OutputT> combineFn) {
       this.combineFn = combineFn;
       this.accumCoder = accumCoder;
     }
@@ -249,6 +358,27 @@ public class StateSpecs {
       return visitor.bindKeyedCombiningValueWithContext(id, this, accumCoder, combineFn);
     }
 
+    @SuppressWarnings("unchecked")
+    @Override
+    public void offerCoders(Coder[] coders) {
+      if (this.accumCoder == null) {
+        if (coders[2] != null) {
+          this.accumCoder = (Coder<AccumT>) coders[2];
+        }
+      }
+    }
+
+    @Override public void finishSpecifying() {
+      if (accumCoder == null) {
+        throw new IllegalStateException("Unable to infer a coder for"
+            + " KeyedCombiningValueWithContextState and no Coder was specified."
+            + " Please set a coder by either invoking"
+            + " StateSpecs.keyedCombiningValue(Coder<AccumT> accumCoder,"
+            + " KeyedCombineFn<K, InputT, AccumT, OutputT> combineFn)"
+            + " or by registering the coder in the Pipeline's CoderRegistry.");
+      }
+    }
+
     @Override
     public boolean equals(Object obj) {
       if (obj == this) {
@@ -282,19 +412,45 @@ public class StateSpecs {
   private static class KeyedCombiningValueStateSpec<K, InputT, AccumT, OutputT>
       implements StateSpec<K, AccumulatorCombiningState<InputT, AccumT, OutputT>> {
 
-    private final Coder<AccumT> accumCoder;
+    @Nullable
+    private Coder<AccumT> accumCoder;
     private final KeyedCombineFn<K, InputT, AccumT, OutputT> keyedCombineFn;
 
     protected KeyedCombiningValueStateSpec(
-        Coder<AccumT> accumCoder, KeyedCombineFn<K, InputT, AccumT, OutputT> keyedCombineFn) {
+        @Nullable Coder<AccumT> accumCoder,
+        KeyedCombineFn<K, InputT, AccumT, OutputT> keyedCombineFn) {
       this.keyedCombineFn = keyedCombineFn;
       this.accumCoder = accumCoder;
     }
 
+    protected Coder<AccumT> getAccumCoder() {
+      return accumCoder;
+    }
+
     @Override
     public AccumulatorCombiningState<InputT, AccumT, OutputT> bind(
         String id, StateBinder<? extends K> visitor) {
-      return visitor.bindKeyedCombiningValue(id, this, accumCoder, keyedCombineFn);
+      return visitor.bindKeyedCombiningValue(id, this, getAccumCoder(), keyedCombineFn);
+    }
+
+    @SuppressWarnings("unchecked")
+    @Override
+    public void offerCoders(Coder[] coders) {
+      if (this.accumCoder == null) {
+        if (coders[2] != null) {
+          this.accumCoder = (Coder<AccumT>) coders[2];
+        }
+      }
+    }
+
+    @Override public void finishSpecifying() {
+      if (getAccumCoder() == null) {
+        throw new IllegalStateException("Unable to infer a coder for CombiningState and no"
+            + " Coder was specified. Please set a coder by either invoking"
+            + " StateSpecs.combiningValue(Coder<AccumT> accumCoder,"
+            + " CombineFn<InputT, AccumT, OutputT> combineFn)"
+            + " or by registering the coder in the Pipeline's CoderRegistry.");
+      }
     }
 
     @Override
@@ -330,9 +486,10 @@ public class StateSpecs {
    */
   private static class BagStateSpec<T> implements StateSpec<Object, BagState<T>> {
 
-    private final Coder<T> elemCoder;
+    @Nullable
+    private Coder<T> elemCoder;
 
-    private BagStateSpec(Coder<T> elemCoder) {
+    private BagStateSpec(@Nullable Coder<T> elemCoder) {
       this.elemCoder = elemCoder;
     }
 
@@ -341,6 +498,25 @@ public class StateSpecs {
       return visitor.bindBag(id, this, elemCoder);
     }
 
+    @SuppressWarnings("unchecked")
+    @Override
+    public void offerCoders(Coder[] coders) {
+      if (this.elemCoder == null) {
+        if (coders[0] != null) {
+          this.elemCoder = (Coder<T>) coders[0];
+        }
+      }
+    }
+
+    @Override public void finishSpecifying() {
+      if (elemCoder == null) {
+        throw new IllegalStateException("Unable to infer a coder for BagState and no Coder"
+            + " was specified. Please set a coder by either invoking"
+            + " StateSpecs.bag(Coder<T> elemCoder) or by registering the coder in the"
+            + " Pipeline's CoderRegistry.");
+      }
+    }
+
     @Override
     public boolean equals(Object obj) {
       if (obj == this) {
@@ -363,10 +539,12 @@ public class StateSpecs {
 
   private static class MapStateSpec<K, V> implements StateSpec<Object, MapState<K, V>> {
 
-    private final Coder<K> keyCoder;
-    private final Coder<V> valueCoder;
+    @Nullable
+    private Coder<K> keyCoder;
+    @Nullable
+    private Coder<V> valueCoder;
 
-    private MapStateSpec(Coder<K> keyCoder, Coder<V> valueCoder) {
+    private MapStateSpec(@Nullable Coder<K> keyCoder, @Nullable Coder<V> valueCoder) {
       this.keyCoder = keyCoder;
       this.valueCoder = valueCoder;
     }
@@ -376,6 +554,30 @@ public class StateSpecs {
       return visitor.bindMap(id, this, keyCoder, valueCoder);
     }
 
+    @SuppressWarnings("unchecked")
+    @Override
+    public void offerCoders(Coder[] coders) {
+      if (this.keyCoder == null) {
+        if (coders[0] != null) {
+          this.keyCoder = (Coder<K>) coders[0];
+        }
+      }
+      if (this.valueCoder == null) {
+        if (coders[1] != null) {
+          this.valueCoder = (Coder<V>) coders[1];
+        }
+      }
+    }
+
+    @Override public void finishSpecifying() {
+      if (keyCoder == null || valueCoder == null) {
+        throw new IllegalStateException("Unable to infer a coder for MapState and no Coder"
+            + " was specified. Please set a coder by either invoking"
+            + " StateSpecs.map(Coder<K> keyCoder, Coder<V> valueCoder) or by registering the"
+            + " coder in the Pipeline's CoderRegistry.");
+      }
+    }
+
     @Override
     public boolean equals(Object obj) {
       if (obj == this) {
@@ -404,9 +606,10 @@ public class StateSpecs {
    */
   private static class SetStateSpec<T> implements StateSpec<Object, SetState<T>> {
 
-    private final Coder<T> elemCoder;
+    @Nullable
+    private Coder<T> elemCoder;
 
-    private SetStateSpec(Coder<T> elemCoder) {
+    private SetStateSpec(@Nullable Coder<T> elemCoder) {
       this.elemCoder = elemCoder;
     }
 
@@ -415,6 +618,25 @@ public class StateSpecs {
       return visitor.bindSet(id, this, elemCoder);
     }
 
+    @SuppressWarnings("unchecked")
+    @Override
+    public void offerCoders(Coder[] coders) {
+      if (this.elemCoder == null) {
+        if (coders[0] != null) {
+          this.elemCoder = (Coder<T>) coders[0];
+        }
+      }
+    }
+
+    @Override public void finishSpecifying() {
+      if (elemCoder == null) {
+        throw new IllegalStateException("Unable to infer a coder for SetState and no Coder"
+            + " was specified. Please set a coder by either invoking"
+            + " StateSpecs.set(Coder<T> elemCoder) or by registering the coder in the"
+            + " Pipeline's CoderRegistry.");
+      }
+    }
+
     @Override
     public boolean equals(Object obj) {
       if (obj == this) {
@@ -461,6 +683,14 @@ public class StateSpecs {
     }
 
     @Override
+    public void offerCoders(Coder[] coders) {
+    }
+
+    @Override public void finishSpecifying() {
+      // Currently an empty implementation as there are no coders to validate.
+    }
+
+    @Override
     public boolean equals(Object obj) {
       if (obj == this) {
         return true;

http://git-wip-us.apache.org/repos/asf/beam/blob/42e690e8/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java
index cbbbe5f..4249a77 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java
@@ -39,6 +39,7 @@ import static org.junit.Assert.assertThat;
 
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.google.common.base.MoreObjects;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Sets;
@@ -55,8 +56,12 @@ import java.util.Map;
 import java.util.Set;
 import org.apache.beam.sdk.Pipeline.PipelineExecutionException;
 import org.apache.beam.sdk.coders.AtomicCoder;
+import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.coders.CoderException;
+import org.apache.beam.sdk.coders.CustomCoder;
 import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.coders.ListCoder;
+import org.apache.beam.sdk.coders.SetCoder;
 import org.apache.beam.sdk.coders.StringUtf8Coder;
 import org.apache.beam.sdk.coders.VarIntCoder;
 import org.apache.beam.sdk.io.CountingInput;
@@ -1036,6 +1041,71 @@ public class ParDoTest implements Serializable {
      }
   }
 
+  private static class MyInteger implements Comparable<MyInteger> {
+    private final int value;
+
+    MyInteger(int value) {
+      this.value = value;
+    }
+
+    public int getValue() {
+      return value;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) {
+        return true;
+      }
+
+      if (!(o instanceof MyInteger)) {
+        return false;
+      }
+
+      MyInteger myInteger = (MyInteger) o;
+
+      return value == myInteger.value;
+
+    }
+
+    @Override
+    public int hashCode() {
+      return value;
+    }
+
+    @Override
+    public int compareTo(MyInteger o) {
+      return Integer.compare(this.getValue(), o.getValue());
+    }
+
+    @Override
+    public String toString() {
+      return "MyInteger{" + "value=" + value + '}';
+    }
+  }
+
+  private static class MyIntegerCoder extends CustomCoder<MyInteger> {
+    private static final MyIntegerCoder INSTANCE = new MyIntegerCoder();
+
+    private final VarIntCoder delegate = VarIntCoder.of();
+
+    public static MyIntegerCoder of() {
+      return INSTANCE;
+    }
+
+    @Override
+    public void encode(MyInteger value, OutputStream outStream, Context context)
+        throws CoderException, IOException {
+      delegate.encode(value.getValue(), outStream, context);
+    }
+
+    @Override
+    public MyInteger decode(InputStream inStream, Context context) throws CoderException,
+        IOException {
+      return new MyInteger(delegate.decode(inStream, context));
+    }
+  }
+
   /** PAssert "matcher" for expected output. */
   static class HasExpectedOutput
       implements SerializableFunction<Iterable<String>, Void>, Serializable {
@@ -1619,6 +1689,132 @@ public class ParDoTest implements Serializable {
 
   @Test
   @Category({ValidatesRunner.class, UsesStatefulParDo.class})
+  public void testValueStateCoderInference() {
+    final String stateId = "foo";
+    MyIntegerCoder myIntegerCoder = MyIntegerCoder.of();
+    pipeline.getCoderRegistry().registerCoder(MyInteger.class, myIntegerCoder);
+
+    DoFn<KV<String, Integer>, MyInteger> fn =
+        new DoFn<KV<String, Integer>, MyInteger>() {
+
+          @StateId(stateId)
+          private final StateSpec<Object, ValueState<MyInteger>> intState =
+              StateSpecs.value();
+
+          @ProcessElement
+          public void processElement(
+              ProcessContext c, @StateId(stateId) ValueState<MyInteger> state) {
+            MyInteger currentValue = MoreObjects.firstNonNull(state.read(), new MyInteger(0));
+            c.output(currentValue);
+            state.write(new MyInteger(currentValue.getValue() + 1));
+          }
+        };
+
+    PCollection<MyInteger> output =
+        pipeline.apply(Create.of(KV.of("hello", 42), KV.of("hello", 97), KV.of("hello", 84)))
+            .apply(ParDo.of(fn)).setCoder(myIntegerCoder);
+
+    PAssert.that(output).containsInAnyOrder(new MyInteger(0), new MyInteger(1), new MyInteger(2));
+    pipeline.run();
+  }
+
+  @Test
+  @Category({ValidatesRunner.class, UsesStatefulParDo.class})
+  public void testValueStateCoderInferenceFailure() throws Exception {
+    final String stateId = "foo";
+    MyIntegerCoder myIntegerCoder = MyIntegerCoder.of();
+
+    DoFn<KV<String, Integer>, MyInteger> fn =
+        new DoFn<KV<String, Integer>, MyInteger>() {
+
+          @StateId(stateId)
+          private final StateSpec<Object, ValueState<MyInteger>> intState =
+              StateSpecs.value();
+
+          @ProcessElement
+          public void processElement(
+              ProcessContext c, @StateId(stateId) ValueState<MyInteger> state) {
+            MyInteger currentValue = MoreObjects.firstNonNull(state.read(), new MyInteger(0));
+            c.output(currentValue);
+            state.write(new MyInteger(currentValue.getValue() + 1));
+          }
+        };
+
+    thrown.expect(RuntimeException.class);
+    thrown.expectMessage("Unable to infer a coder for ValueState and no Coder was specified.");
+
+    pipeline.apply(Create.of(KV.of("hello", 42), KV.of("hello", 97), KV.of("hello", 84)))
+        .apply(ParDo.of(fn)).setCoder(myIntegerCoder);
+
+    pipeline.run();
+  }
+
+  @Test
+  @Category({ValidatesRunner.class, UsesStatefulParDo.class})
+  public void testValueStateCoderInferenceFromInputCoder() {
+    final String stateId = "foo";
+    MyIntegerCoder myIntegerCoder = MyIntegerCoder.of();
+
+    DoFn<KV<String, MyInteger>, MyInteger> fn =
+        new DoFn<KV<String, MyInteger>, MyInteger>() {
+
+          @StateId(stateId)
+          private final StateSpec<Object, ValueState<MyInteger>> intState =
+              StateSpecs.value();
+
+          @ProcessElement
+          public void processElement(
+              ProcessContext c, @StateId(stateId) ValueState<MyInteger> state) {
+            MyInteger currentValue = MoreObjects.firstNonNull(state.read(), new MyInteger(0));
+            c.output(currentValue);
+            state.write(new MyInteger(currentValue.getValue() + 1));
+          }
+        };
+
+        pipeline
+            .apply(Create.of(KV.of("hello", new MyInteger(42)),
+                KV.of("hello", new MyInteger(97)), KV.of("hello", new MyInteger(84)))
+                .withCoder(KvCoder.of(StringUtf8Coder.of(), myIntegerCoder)))
+            .apply(ParDo.of(fn)).setCoder(myIntegerCoder);
+
+    pipeline.run();
+  }
+
+  @Test
+  @Category({ValidatesRunner.class, UsesStatefulParDo.class})
+  public void testCoderInferenceOfList() {
+    final String stateId = "foo";
+    MyIntegerCoder myIntegerCoder = MyIntegerCoder.of();
+    pipeline.getCoderRegistry().registerCoder(MyInteger.class, myIntegerCoder);
+
+    DoFn<KV<String, Integer>, List<MyInteger>> fn =
+        new DoFn<KV<String, Integer>, List<MyInteger>>() {
+
+          @StateId(stateId)
+          private final StateSpec<Object, ValueState<List<MyInteger>>> intState =
+              StateSpecs.value();
+
+          @ProcessElement
+          public void processElement(
+              ProcessContext c, @StateId(stateId) ValueState<List<MyInteger>> state) {
+            MyInteger myInteger = new MyInteger(c.element().getValue());
+            List<MyInteger> currentValue = state.read();
+            List<MyInteger> newValue = currentValue != null
+                ? ImmutableList.<MyInteger>builder().addAll(currentValue).add(myInteger).build()
+                : Collections.singletonList(myInteger);
+            c.output(newValue);
+            state.write(newValue);
+          }
+        };
+
+    pipeline.apply(Create.of(KV.of("hello", 42), KV.of("hello", 97), KV.of("hello", 84)))
+        .apply(ParDo.of(fn)).setCoder(ListCoder.of(myIntegerCoder));
+
+    pipeline.run();
+  }
+
+  @Test
+  @Category({ValidatesRunner.class, UsesStatefulParDo.class})
   public void testValueStateFixedWindows() {
     final String stateId = "foo";
 
@@ -1801,6 +1997,82 @@ public class ParDoTest implements Serializable {
   }
 
   @Test
+  @Category({ValidatesRunner.class, UsesStatefulParDo.class})
+  public void testBagStateCoderInference() {
+    final String stateId = "foo";
+    Coder<MyInteger> myIntegerCoder = MyIntegerCoder.of();
+    pipeline.getCoderRegistry().registerCoder(MyInteger.class, myIntegerCoder);
+
+    DoFn<KV<String, Integer>, List<MyInteger>> fn =
+        new DoFn<KV<String, Integer>, List<MyInteger>>() {
+
+          @StateId(stateId)
+          private final StateSpec<Object, BagState<MyInteger>> bufferState =
+              StateSpecs.bag();
+
+          @ProcessElement
+          public void processElement(
+              ProcessContext c, @StateId(stateId) BagState<MyInteger> state) {
+            Iterable<MyInteger> currentValue = state.read();
+            state.add(new MyInteger(c.element().getValue()));
+            if (Iterables.size(state.read()) >= 4) {
+              List<MyInteger> sorted = Lists.newArrayList(currentValue);
+              Collections.sort(sorted);
+              c.output(sorted);
+            }
+          }
+        };
+
+    PCollection<List<MyInteger>> output =
+        pipeline.apply(
+            Create.of(
+                KV.of("hello", 97), KV.of("hello", 42), KV.of("hello", 84), KV.of("hello", 12)))
+            .apply(ParDo.of(fn)).setCoder(ListCoder.of(myIntegerCoder));
+
+    PAssert.that(output).containsInAnyOrder(Lists.newArrayList(new MyInteger(12), new MyInteger(42),
+        new MyInteger(84), new MyInteger(97)));
+
+    pipeline.run();
+  }
+
+  @Test
+  @Category({ValidatesRunner.class, UsesStatefulParDo.class})
+  public void testBagStateCoderInferenceFailure() throws Exception {
+    final String stateId = "foo";
+    Coder<MyInteger> myIntegerCoder = MyIntegerCoder.of();
+
+    DoFn<KV<String, Integer>, List<MyInteger>> fn =
+        new DoFn<KV<String, Integer>, List<MyInteger>>() {
+
+          @StateId(stateId)
+          private final StateSpec<Object, BagState<MyInteger>> bufferState =
+              StateSpecs.bag();
+
+          @ProcessElement
+          public void processElement(
+              ProcessContext c, @StateId(stateId) BagState<MyInteger> state) {
+            Iterable<MyInteger> currentValue = state.read();
+            state.add(new MyInteger(c.element().getValue()));
+            if (Iterables.size(state.read()) >= 4) {
+              List<MyInteger> sorted = Lists.newArrayList(currentValue);
+              Collections.sort(sorted);
+              c.output(sorted);
+            }
+          }
+        };
+
+    thrown.expect(RuntimeException.class);
+    thrown.expectMessage("Unable to infer a coder for BagState and no Coder was specified.");
+
+    pipeline.apply(
+        Create.of(
+            KV.of("hello", 97), KV.of("hello", 42), KV.of("hello", 84), KV.of("hello", 12)))
+        .apply(ParDo.of(fn)).setCoder(ListCoder.of(myIntegerCoder));
+
+    pipeline.run();
+  }
+
+  @Test
   @Category({ValidatesRunner.class, UsesStatefulParDo.class, UsesSetState.class})
   public void testSetState() {
     final String stateId = "foo";
@@ -1843,6 +2115,93 @@ public class ParDoTest implements Serializable {
   }
 
   @Test
+  @Category({ValidatesRunner.class, UsesStatefulParDo.class, UsesSetState.class})
+  public void testSetStateCoderInference() {
+    final String stateId = "foo";
+    final String countStateId = "count";
+    Coder<MyInteger> myIntegerCoder = MyIntegerCoder.of();
+    pipeline.getCoderRegistry().registerCoder(MyInteger.class, myIntegerCoder);
+
+    DoFn<KV<String, Integer>, Set<MyInteger>> fn =
+        new DoFn<KV<String, Integer>, Set<MyInteger>>() {
+
+          @StateId(stateId)
+          private final StateSpec<Object, SetState<MyInteger>> setState = StateSpecs.set();
+
+          @StateId(countStateId)
+          private final StateSpec<Object, AccumulatorCombiningState<Integer, int[], Integer>>
+              countState = StateSpecs.combiningValueFromInputInternal(VarIntCoder.of(),
+              Sum.ofIntegers());
+
+          @ProcessElement
+          public void processElement(
+              ProcessContext c,
+              @StateId(stateId) SetState<MyInteger> state,
+              @StateId(countStateId) AccumulatorCombiningState<Integer, int[], Integer> count) {
+            state.add(new MyInteger(c.element().getValue()));
+            count.add(1);
+            if (count.read() >= 4) {
+              Set<MyInteger> set = Sets.newHashSet(state.read());
+              c.output(set);
+            }
+          }
+        };
+
+    PCollection<Set<MyInteger>> output =
+        pipeline.apply(
+            Create.of(
+                KV.of("hello", 97), KV.of("hello", 42), KV.of("hello", 42), KV.of("hello", 12)))
+            .apply(ParDo.of(fn)).setCoder(SetCoder.of(myIntegerCoder));
+
+    PAssert.that(output).containsInAnyOrder(
+        Sets.newHashSet(new MyInteger(97), new MyInteger(42), new MyInteger(12)));
+    pipeline.run();
+  }
+
+  @Test
+  @Category({ValidatesRunner.class, UsesStatefulParDo.class, UsesSetState.class})
+  public void testSetStateCoderInferenceFailure() throws Exception {
+    final String stateId = "foo";
+    final String countStateId = "count";
+    Coder<MyInteger> myIntegerCoder = MyIntegerCoder.of();
+
+    DoFn<KV<String, Integer>, Set<MyInteger>> fn =
+        new DoFn<KV<String, Integer>, Set<MyInteger>>() {
+
+          @StateId(stateId)
+          private final StateSpec<Object, SetState<MyInteger>> setState = StateSpecs.set();
+
+          @StateId(countStateId)
+          private final StateSpec<Object, AccumulatorCombiningState<Integer, int[], Integer>>
+              countState = StateSpecs.combiningValueFromInputInternal(VarIntCoder.of(),
+              Sum.ofIntegers());
+
+          @ProcessElement
+          public void processElement(
+              ProcessContext c,
+              @StateId(stateId) SetState<MyInteger> state,
+              @StateId(countStateId) AccumulatorCombiningState<Integer, int[], Integer> count) {
+            state.add(new MyInteger(c.element().getValue()));
+            count.add(1);
+            if (count.read() >= 4) {
+              Set<MyInteger> set = Sets.newHashSet(state.read());
+              c.output(set);
+            }
+          }
+        };
+
+    thrown.expect(RuntimeException.class);
+    thrown.expectMessage("Unable to infer a coder for SetState and no Coder was specified.");
+
+    pipeline.apply(
+        Create.of(
+            KV.of("hello", 97), KV.of("hello", 42), KV.of("hello", 42), KV.of("hello", 12)))
+        .apply(ParDo.of(fn)).setCoder(SetCoder.of(myIntegerCoder));
+
+    pipeline.run();
+  }
+
+  @Test
   @Category({ValidatesRunner.class, UsesStatefulParDo.class, UsesMapState.class})
   public void testMapState() {
     final String stateId = "foo";
@@ -1888,6 +2247,99 @@ public class ParDoTest implements Serializable {
   }
 
   @Test
+  @Category({ValidatesRunner.class, UsesStatefulParDo.class, UsesMapState.class})
+  public void testMapStateCoderInference() {
+    final String stateId = "foo";
+    final String countStateId = "count";
+    Coder<MyInteger> myIntegerCoder = MyIntegerCoder.of();
+    pipeline.getCoderRegistry().registerCoder(MyInteger.class, myIntegerCoder);
+
+    DoFn<KV<String, KV<String, Integer>>, KV<String, MyInteger>> fn =
+        new DoFn<KV<String, KV<String, Integer>>, KV<String, MyInteger>>() {
+
+          @StateId(stateId)
+          private final StateSpec<Object, MapState<String, MyInteger>> mapState = StateSpecs.map();
+          @StateId(countStateId)
+          private final StateSpec<Object, AccumulatorCombiningState<Integer, int[], Integer>>
+              countState = StateSpecs.combiningValueFromInputInternal(VarIntCoder.of(),
+              Sum.ofIntegers());
+
+          @ProcessElement
+          public void processElement(
+              ProcessContext c, @StateId(stateId) MapState<String, MyInteger> state,
+              @StateId(countStateId) AccumulatorCombiningState<Integer, int[], Integer>
+                  count) {
+            KV<String, Integer> value = c.element().getValue();
+            state.put(value.getKey(), new MyInteger(value.getValue()));
+            count.add(1);
+            if (count.read() >= 4) {
+              Iterable<Map.Entry<String, MyInteger>> iterate = state.iterate();
+              for (Map.Entry<String, MyInteger> entry : iterate) {
+                c.output(KV.of(entry.getKey(), entry.getValue()));
+              }
+            }
+          }
+        };
+
+    PCollection<KV<String, MyInteger>> output =
+        pipeline.apply(
+            Create.of(
+                KV.of("hello", KV.of("a", 97)), KV.of("hello", KV.of("b", 42)),
+                KV.of("hello", KV.of("b", 42)), KV.of("hello", KV.of("c", 12))))
+            .apply(ParDo.of(fn)).setCoder(KvCoder.of(StringUtf8Coder.of(), myIntegerCoder));
+
+    PAssert.that(output).containsInAnyOrder(KV.of("a", new MyInteger(97)),
+        KV.of("b", new MyInteger(42)), KV.of("c", new MyInteger(12)));
+    pipeline.run();
+  }
+
+  @Test
+  @Category({ValidatesRunner.class, UsesStatefulParDo.class, UsesMapState.class})
+  public void testMapStateCoderInferenceFailure() throws Exception {
+    final String stateId = "foo";
+    final String countStateId = "count";
+    Coder<MyInteger> myIntegerCoder = MyIntegerCoder.of();
+
+    DoFn<KV<String, KV<String, Integer>>, KV<String, MyInteger>> fn =
+        new DoFn<KV<String, KV<String, Integer>>, KV<String, MyInteger>>() {
+
+          @StateId(stateId)
+          private final StateSpec<Object, MapState<String, MyInteger>> mapState = StateSpecs.map();
+          @StateId(countStateId)
+          private final StateSpec<Object, AccumulatorCombiningState<Integer, int[], Integer>>
+              countState = StateSpecs.combiningValueFromInputInternal(VarIntCoder.of(),
+              Sum.ofIntegers());
+
+          @ProcessElement
+          public void processElement(
+              ProcessContext c, @StateId(stateId) MapState<String, MyInteger> state,
+              @StateId(countStateId) AccumulatorCombiningState<Integer, int[], Integer>
+                  count) {
+            KV<String, Integer> value = c.element().getValue();
+            state.put(value.getKey(), new MyInteger(value.getValue()));
+            count.add(1);
+            if (count.read() >= 4) {
+              Iterable<Map.Entry<String, MyInteger>> iterate = state.iterate();
+              for (Map.Entry<String, MyInteger> entry : iterate) {
+                c.output(KV.of(entry.getKey(), entry.getValue()));
+              }
+            }
+          }
+        };
+
+    thrown.expect(RuntimeException.class);
+    thrown.expectMessage("Unable to infer a coder for MapState and no Coder was specified.");
+
+    pipeline.apply(
+        Create.of(
+            KV.of("hello", KV.of("a", 97)), KV.of("hello", KV.of("b", 42)),
+            KV.of("hello", KV.of("b", 42)), KV.of("hello", KV.of("c", 12))))
+        .apply(ParDo.of(fn)).setCoder(KvCoder.of(StringUtf8Coder.of(), myIntegerCoder));
+
+    pipeline.run();
+  }
+
+  @Test
   @Category({ValidatesRunner.class, UsesStatefulParDo.class})
   public void testCombiningState() {
     final String stateId = "foo";
@@ -1928,6 +2380,132 @@ public class ParDoTest implements Serializable {
 
   @Test
   @Category({ValidatesRunner.class, UsesStatefulParDo.class})
+  public void testCombiningStateCoderInference() {
+    pipeline.getCoderRegistry().registerCoder(MyInteger.class, MyIntegerCoder.of());
+
+    final String stateId = "foo";
+
+    DoFn<KV<String, Integer>, String> fn =
+        new DoFn<KV<String, Integer>, String>() {
+          private static final int EXPECTED_SUM = 16;
+
+          @StateId(stateId)
+          private final StateSpec<
+              Object, AccumulatorCombiningState<Integer, MyInteger, Integer>>
+              combiningState =
+              StateSpecs.combiningValue(new Combine.CombineFn<Integer, MyInteger, Integer>() {
+                @Override
+                public MyInteger createAccumulator() {
+                  return new MyInteger(0);
+                }
+
+                @Override
+                public MyInteger addInput(MyInteger accumulator, Integer input) {
+                  return new MyInteger(accumulator.getValue() + input);
+                }
+
+                @Override
+                public MyInteger mergeAccumulators(Iterable<MyInteger> accumulators) {
+                  int newValue = 0;
+                  for (MyInteger myInteger : accumulators) {
+                    newValue += myInteger.getValue();
+                  }
+                  return new MyInteger(newValue);
+                }
+
+                @Override
+                public Integer extractOutput(MyInteger accumulator) {
+                  return accumulator.getValue();
+                }
+              });
+
+          @ProcessElement
+          public void processElement(
+              ProcessContext c,
+              @StateId(stateId)
+                  AccumulatorCombiningState<Integer, MyInteger, Integer> state) {
+            state.add(c.element().getValue());
+            Integer currentValue = state.read();
+            if (currentValue == EXPECTED_SUM) {
+              c.output("right on");
+            }
+          }
+        };
+
+    PCollection<String> output =
+        pipeline
+            .apply(Create.of(KV.of("hello", 3), KV.of("hello", 6), KV.of("hello", 7)))
+            .apply(ParDo.of(fn));
+
+    // There should only be one moment at which the average is exactly 16
+    PAssert.that(output).containsInAnyOrder("right on");
+    pipeline.run();
+  }
+
+  @Test
+  @Category({ValidatesRunner.class, UsesStatefulParDo.class})
+  public void testCombiningStateCoderInferenceFailure() throws Exception {
+    final String stateId = "foo";
+
+    DoFn<KV<String, Integer>, String> fn =
+        new DoFn<KV<String, Integer>, String>() {
+          private static final int EXPECTED_SUM = 16;
+
+          @StateId(stateId)
+          private final StateSpec<
+              Object, AccumulatorCombiningState<Integer, MyInteger, Integer>>
+              combiningState =
+              StateSpecs.combiningValue(new Combine.CombineFn<Integer, MyInteger, Integer>() {
+                @Override
+                public MyInteger createAccumulator() {
+                  return new MyInteger(0);
+                }
+
+                @Override
+                public MyInteger addInput(MyInteger accumulator, Integer input) {
+                  return new MyInteger(accumulator.getValue() + input);
+                }
+
+                @Override
+                public MyInteger mergeAccumulators(Iterable<MyInteger> accumulators) {
+                  int newValue = 0;
+                  for (MyInteger myInteger : accumulators) {
+                    newValue += myInteger.getValue();
+                  }
+                  return new MyInteger(newValue);
+                }
+
+                @Override
+                public Integer extractOutput(MyInteger accumulator) {
+                  return accumulator.getValue();
+                }
+              });
+
+          @ProcessElement
+          public void processElement(
+              ProcessContext c,
+              @StateId(stateId)
+                  AccumulatorCombiningState<Integer, MyInteger, Integer> state) {
+            state.add(c.element().getValue());
+            Integer currentValue = state.read();
+            if (currentValue == EXPECTED_SUM) {
+              c.output("right on");
+            }
+          }
+        };
+
+    thrown.expect(RuntimeException.class);
+    thrown.expectMessage("Unable to infer a coder for CombiningState and no Coder was specified.");
+
+    pipeline
+        .apply(Create.of(KV.of("hello", 3), KV.of("hello", 6), KV.of("hello", 7)))
+        .apply(ParDo.of(fn));
+
+    pipeline.run();
+  }
+
+  @Test
+  @Category({ValidatesRunner.class, UsesStatefulParDo.class})
   public void testBagStateSideInput() {
 
     final PCollectionView<List<Integer>> listView =


Mime
View raw message