This closes #3233: [BEAM-115] Runner API Translations for StateSpec and TimerSpec

  Implement TimerSpec and StateSpec translation
  Make Java serialized CombineFn URN public
  Add case dispatch to StateSpec
  Flesh out TimerSpec and StateSpec in Runner API
  Allow translation to throw IOException
  Mark CombineFnWithContext StateSpecs internal

Author: kennknowles

runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/CombineTranslation.java

@@ -49,7 +49,7 @@
  * RunnerApi.CombinePayload} protos.
  */
 public class CombineTranslation {
-  private static final String JAVA_SERIALIZED_COMBINE_FN_URN = "urn:beam:java:combinefn:v1";
+  public static final String JAVA_SERIALIZED_COMBINE_FN_URN = "urn:beam:java:combinefn:v1";
 
   public static CombinePayload toProto(
       AppliedPTransform<?, ?, Combine.PerKey<?, ?, ?>> combine, SdkComponents sdkComponents)
@@ -86,7 +86,7 @@ private static <K, InputT, AccumT> Coder<AccumT> extractAccumulatorCoder(
         .getAccumulatorCoder();
   }
 
-  private static SdkFunctionSpec toProto(GlobalCombineFn<?, ?, ?> combineFn) {
+  public static SdkFunctionSpec toProto(GlobalCombineFn<?, ?, ?> combineFn) {
     return SdkFunctionSpec.newBuilder()
         // TODO: Set Java SDK Environment URN
         .setSpec(

runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PTransformTranslation.java

@@ -138,7 +138,8 @@ public static String urnForTransform(PTransform<?, ?> transform) {
    */
   public interface TransformPayloadTranslator<T extends PTransform<?, ?>> {
     String getUrn(T transform);
-    FunctionSpec translate(AppliedPTransform<?, ?, T> application, SdkComponents components);
+    FunctionSpec translate(AppliedPTransform<?, ?, T> application, SdkComponents components)
+        throws IOException;
   }
 
   /**

runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ParDoTranslation.java

@@ -19,10 +19,12 @@
 package org.apache.beam.runners.core.construction;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Preconditions.checkState;
 import static org.apache.beam.runners.core.construction.PTransformTranslation.PAR_DO_TRANSFORM_URN;
 
 import com.google.auto.service.AutoService;
 import com.google.auto.value.AutoValue;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Optional;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Sets;
@@ -46,9 +48,12 @@
 import org.apache.beam.sdk.common.runner.v1.RunnerApi.SdkFunctionSpec;
 import org.apache.beam.sdk.common.runner.v1.RunnerApi.SideInput;
 import org.apache.beam.sdk.common.runner.v1.RunnerApi.SideInput.Builder;
-import org.apache.beam.sdk.common.runner.v1.RunnerApi.StateSpec;
-import org.apache.beam.sdk.common.runner.v1.RunnerApi.TimerSpec;
 import org.apache.beam.sdk.runners.AppliedPTransform;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.transforms.Combine;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.Materializations;
 import org.apache.beam.sdk.transforms.PTransform;
@@ -107,7 +112,8 @@ public String getUrn(ParDo.MultiOutput<?, ?> transform) {
 
     @Override
     public FunctionSpec translate(
-        AppliedPTransform<?, ?, MultiOutput<?, ?>> transform, SdkComponents components) {
+        AppliedPTransform<?, ?, MultiOutput<?, ?>> transform, SdkComponents components)
+        throws IOException {
       ParDoPayload payload = toProto(transform.getTransform(), components);
       return RunnerApi.FunctionSpec.newBuilder()
           .setUrn(PAR_DO_TRANSFORM_URN)
@@ -128,8 +134,10 @@ public static class Registrar implements TransformPayloadTranslatorRegistrar {
     }
   }
 
-  public static ParDoPayload toProto(ParDo.MultiOutput<?, ?> parDo, SdkComponents components) {
-    DoFnSignature signature = DoFnSignatures.getSignature(parDo.getFn().getClass());
+  public static ParDoPayload toProto(ParDo.MultiOutput<?, ?> parDo, SdkComponents components)
+  throws IOException {
+    DoFn<?, ?> doFn = parDo.getFn();
+    DoFnSignature signature = DoFnSignatures.getSignature(doFn.getClass());
     Map<String, StateDeclaration> states = signature.stateDeclarations();
     Map<String, TimerDeclaration> timers = signature.timerDeclarations();
     List<Parameter> parameters = signature.processElement().extraParameters();
@@ -146,16 +154,62 @@ public static ParDoPayload toProto(ParDo.MultiOutput<?, ?> parDo, SdkComponents
       }
     }
     for (Map.Entry<String, StateDeclaration> state : states.entrySet()) {
-      StateSpec spec = toProto(state.getValue());
+      RunnerApi.StateSpec spec =
+          toProto(getStateSpecOrCrash(state.getValue(), doFn), components);
       builder.putStateSpecs(state.getKey(), spec);
     }
     for (Map.Entry<String, TimerDeclaration> timer : timers.entrySet()) {
-      TimerSpec spec = toProto(timer.getValue());
+      RunnerApi.TimerSpec spec =
+          toProto(getTimerSpecOrCrash(timer.getValue(), doFn));
       builder.putTimerSpecs(timer.getKey(), spec);
     }
     return builder.build();
   }
 
+  private static StateSpec<?> getStateSpecOrCrash(
+      StateDeclaration stateDeclaration, DoFn<?, ?> target) {
+    try {
+      Object fieldValue = stateDeclaration.field().get(target);
+      checkState(fieldValue instanceof StateSpec,
+          "Malformed %s class %s: state declaration field %s does not have type %s.",
+          DoFn.class.getSimpleName(),
+          target.getClass().getName(),
+          stateDeclaration.field().getName(),
+          StateSpec.class);
+
+      return (StateSpec<?>) stateDeclaration.field().get(target);
+    } catch (IllegalAccessException exc) {
+      throw new RuntimeException(
+          String.format(
+              "Malformed %s class %s: state declaration field %s is not accessible.",
+              DoFn.class.getSimpleName(),
+              target.getClass().getName(),
+              stateDeclaration.field().getName()));
+    }
+  }
+
+  private static TimerSpec getTimerSpecOrCrash(
+      TimerDeclaration timerDeclaration, DoFn<?, ?> target) {
+    try {
+      Object fieldValue = timerDeclaration.field().get(target);
+      checkState(fieldValue instanceof TimerSpec,
+          "Malformed %s class %s: timer declaration field %s does not have type %s.",
+          DoFn.class.getSimpleName(),
+          target.getClass().getName(),
+          timerDeclaration.field().getName(),
+          TimerSpec.class);
+
+      return (TimerSpec) timerDeclaration.field().get(target);
+    } catch (IllegalAccessException exc) {
+      throw new RuntimeException(
+          String.format(
+              "Malformed %s class %s: timer declaration field %s is not accessible.",
+              DoFn.class.getSimpleName(),
+              target.getClass().getName(),
+              timerDeclaration.field().getName()));
+    }
+  }
+
   public static DoFn<?, ?> getDoFn(ParDoPayload payload) throws InvalidProtocolBufferException {
     return doFnAndMainOutputTagFromProto(payload.getDoFn()).getDoFn();
   }
@@ -179,14 +233,149 @@ public static RunnerApi.PCollection getMainInput(
     return components.getPcollectionsOrThrow(ptransform.getInputsOrThrow(mainInputId));
   }
 
-  // TODO: Implement
-  private static StateSpec toProto(StateDeclaration state) {
-    throw new UnsupportedOperationException("Not yet supported");
+  @VisibleForTesting
+  static RunnerApi.StateSpec toProto(StateSpec<?> stateSpec, final SdkComponents components)
+      throws IOException {
+    final RunnerApi.StateSpec.Builder builder = RunnerApi.StateSpec.newBuilder();
+
+    return stateSpec.match(
+        new StateSpec.Cases<RunnerApi.StateSpec>() {
+          @Override
+          public RunnerApi.StateSpec dispatchValue(Coder<?> valueCoder) {
+            return builder
+                .setValueSpec(
+                    RunnerApi.ValueStateSpec.newBuilder()
+                        .setCoderId(registerCoderOrThrow(components, valueCoder)))
+                .build();
+          }
+
+          @Override
+          public RunnerApi.StateSpec dispatchBag(Coder<?> elementCoder) {
+            return builder
+                .setBagSpec(
+                    RunnerApi.BagStateSpec.newBuilder()
+                        .setElementCoderId(registerCoderOrThrow(components, elementCoder)))
+                .build();
+          }
+
+          @Override
+          public RunnerApi.StateSpec dispatchCombining(
+              Combine.CombineFn<?, ?, ?> combineFn, Coder<?> accumCoder) {
+            return builder
+                .setCombiningSpec(
+                    RunnerApi.CombiningStateSpec.newBuilder()
+                        .setAccumulatorCoderId(registerCoderOrThrow(components, accumCoder))
+                        .setCombineFn(CombineTranslation.toProto(combineFn)))
+                .build();
+          }
+
+          @Override
+          public RunnerApi.StateSpec dispatchMap(Coder<?> keyCoder, Coder<?> valueCoder) {
+            return builder
+                .setMapSpec(
+                    RunnerApi.MapStateSpec.newBuilder()
+                        .setKeyCoderId(registerCoderOrThrow(components, keyCoder))
+                        .setValueCoderId(registerCoderOrThrow(components, valueCoder)))
+                .build();
+          }
+
+          @Override
+          public RunnerApi.StateSpec dispatchSet(Coder<?> elementCoder) {
+            return builder
+                .setSetSpec(
+                    RunnerApi.SetStateSpec.newBuilder()
+                        .setElementCoderId(registerCoderOrThrow(components, elementCoder)))
+                .build();
+          }
+        });
+  }
+
+  @VisibleForTesting
+  static StateSpec<?> fromProto(RunnerApi.StateSpec stateSpec, RunnerApi.Components components)
+      throws IOException {
+    switch (stateSpec.getSpecCase()) {
+      case VALUE_SPEC:
+        return StateSpecs.value(
+            CoderTranslation.fromProto(
+                components.getCodersMap().get(stateSpec.getValueSpec().getCoderId()), components));
+      case BAG_SPEC:
+        return StateSpecs.bag(
+            CoderTranslation.fromProto(
+                components.getCodersMap().get(stateSpec.getBagSpec().getElementCoderId()),
+                components));
+      case COMBINING_SPEC:
+        FunctionSpec combineFnSpec = stateSpec.getCombiningSpec().getCombineFn().getSpec();
+
+        if (!combineFnSpec.getUrn().equals(CombineTranslation.JAVA_SERIALIZED_COMBINE_FN_URN)) {
+          throw new UnsupportedOperationException(
+              String.format(
+                  "Cannot create %s from non-Java %s: %s",
+                  StateSpec.class.getSimpleName(),
+                  Combine.CombineFn.class.getSimpleName(),
+                  combineFnSpec.getUrn()));
+        }
+
+        Combine.CombineFn<?, ?, ?> combineFn =
+            (Combine.CombineFn<?, ?, ?>)
+                SerializableUtils.deserializeFromByteArray(
+                    combineFnSpec.getParameter().unpack(BytesValue.class).toByteArray(),
+                    Combine.CombineFn.class.getSimpleName());
+
+        // Rawtype coder cast because it is required to be a valid accumulator coder
+        // for the CombineFn, by construction
+        return StateSpecs.combining(
+            (Coder)
+                CoderTranslation.fromProto(
+                    components
+                        .getCodersMap()
+                        .get(stateSpec.getCombiningSpec().getAccumulatorCoderId()),
+                    components),
+            combineFn);
+
+      case MAP_SPEC:
+        return StateSpecs.map(
+            CoderTranslation.fromProto(
+                components.getCodersOrThrow(stateSpec.getMapSpec().getKeyCoderId()), components),
+            CoderTranslation.fromProto(
+                components.getCodersOrThrow(stateSpec.getMapSpec().getValueCoderId()), components));
+
+      case SET_SPEC:
+        return StateSpecs.set(
+            CoderTranslation.fromProto(
+                components.getCodersMap().get(stateSpec.getSetSpec().getElementCoderId()),
+                components));
+
+      case SPEC_NOT_SET:
+      default:
+        throw new IllegalArgumentException(
+            String.format("Unknown %s: %s", RunnerApi.StateSpec.class.getName(), stateSpec));
+
+    }
+  }
+
+  private static String registerCoderOrThrow(SdkComponents components, Coder coder) {
+    try {
+      return components.registerCoder(coder);
+    } catch (IOException exc) {
+      throw new RuntimeException("Failure to register coder", exc);
+    }
   }
 
-  // TODO: Implement
-  private static TimerSpec toProto(TimerDeclaration timer) {
-    throw new UnsupportedOperationException("Not yet supported");
+  private static RunnerApi.TimerSpec toProto(TimerSpec timer) {
+    return RunnerApi.TimerSpec.newBuilder().setTimeDomain(toProto(timer.getTimeDomain())).build();
+  }
+
+  private static RunnerApi.TimeDomain toProto(TimeDomain timeDomain) {
+    switch(timeDomain) {
+      case EVENT_TIME:
+        return RunnerApi.TimeDomain.EVENT_TIME;
+      case PROCESSING_TIME:
+        return RunnerApi.TimeDomain.PROCESSING_TIME;
+      case SYNCHRONIZED_PROCESSING_TIME:
+        return RunnerApi.TimeDomain.SYNCHRONIZED_PROCESSING_TIME;
+      default:
+        throw new IllegalArgumentException("Unknown time domain");
+    }
   }
 
   @AutoValue