Merge pull request #11667: [BEAM-9947] Store updated ParDoPayload for length-prefixed timer coders

mxm · web-flow · commit 34dabe3611c9 · 2020-05-12T11:42:51.000+02:00
Backport of #11658.
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ModelCoders.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ModelCoders.java
@@ -115,7 +115,11 @@ public abstract static class WindowedValueCoderComponents {
   }
 
   public static KvCoderComponents getKvCoderComponents(Coder coder) {
-    checkArgument(KV_CODER_URN.equals(coder.getSpec().getUrn()));
+    checkArgument(
+        KV_CODER_URN.equals(coder.getSpec().getUrn()),
+        "Provided coder %s is not of type %s",
+        coder.getSpec().getUrn(),
+        KV_CODER_URN);
     return new AutoValue_ModelCoders_KvCoderComponents(
         coder.getComponentCoderIds(0), coder.getComponentCoderIds(1));
   }
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/ProcessBundleDescriptors.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/ProcessBundleDescriptors.java
@@ -396,18 +396,16 @@ private static Map<String, Map<String, TimerSpec>> forTimerSpecs(
               .getTransformsOrThrow(timerReference.transform().getId())
               .toBuilder()
               .getSpecBuilder();
-      RunnerApi.ParDoPayload updatedPayload =
-          RunnerApi.ParDoPayload.parseFrom(updatedSpec.getPayload());
-      updatedPayload
-          .toBuilder()
-          .putTimerFamilySpecs(
-              timerReference.localName(),
-              updatedPayload
-                  .getTimerFamilySpecsOrThrow(timerReference.localName())
-                  .toBuilder()
-                  .setTimerFamilyCoderId(sdkCoderId)
-                  .build());
-      updatedSpec.setPayload(updatedPayload.toByteString());
+      RunnerApi.ParDoPayload.Builder updatedPayload =
+          RunnerApi.ParDoPayload.parseFrom(updatedSpec.getPayload()).toBuilder();
+      updatedPayload.putTimerFamilySpecs(
+          timerReference.localName(),
+          updatedPayload
+              .getTimerFamilySpecsOrThrow(timerReference.localName())
+              .toBuilder()
+              .setTimerFamilyCoderId(sdkCoderId)
+              .build());
+      updatedSpec.setPayload(updatedPayload.build().toByteString());
       components.putTransforms(
           timerReference.transform().getId(),
           // Since a transform can have more then one timer, update the transform inside components
diff --git a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/ProcessBundleDescriptorsTest.java b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/ProcessBundleDescriptorsTest.java
@@ -34,6 +34,7 @@
 import org.apache.beam.runners.core.construction.graph.FusedPipeline;
 import org.apache.beam.runners.core.construction.graph.GreedyPipelineFuser;
 import org.apache.beam.runners.core.construction.graph.PipelineNode;
+import org.apache.beam.runners.core.construction.graph.TimerReference;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.coders.KvCoder;
@@ -42,6 +43,10 @@
 import org.apache.beam.sdk.state.BagState;
 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.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.GroupByKey;
 import org.apache.beam.sdk.transforms.Impulse;
@@ -59,7 +64,7 @@ public class ProcessBundleDescriptorsTest implements Serializable {
    * LengthPrefixCoder.
    */
   @Test
-  public void testWrapKeyCoderOfStatefulExecutableStageInLengthPrefixCoder() throws Exception {
+  public void testLengthPrefixingOfKeyCoderInStatefulExecutableStage() throws Exception {
     // Add another stateful stage with a non-standard key coder
     Pipeline p = Pipeline.create();
     Coder<Void> keycoder = VoidCoder.of();
@@ -82,16 +87,18 @@ public void process(ProcessContext ctxt) {}
                   private final StateSpec<BagState<String>> bufferState =
                       StateSpecs.bag(StringUtf8Coder.of());
 
+                  @TimerId("timerId")
+                  private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.EVENT_TIME);
+
                   @ProcessElement
                   public void processElement(
                       @Element KV<Void, String> element,
                       @StateId("stateId") BagState<String> state,
-                      OutputReceiver<KV<Void, String>> r) {
-                    for (String value : state.read()) {
-                      r.output(KV.of(element.getKey(), value));
-                    }
-                    state.add(element.getValue());
-                  }
+                      @TimerId("timerId") Timer timer,
+                      OutputReceiver<KV<Void, String>> r) {}
+
+                  @OnTimer("timerId")
+                  public void onTimer() {}
                 }))
         // Force the output to be materialized
         .apply("gbk", GroupByKey.create());
@@ -111,34 +118,46 @@ public void processElement(
 
     // Ensure original key coder is not a LengthPrefixCoder
     Map<String, RunnerApi.Coder> stageCoderMap = stage.getComponents().getCodersMap();
-    RunnerApi.Coder originalCoder =
+    RunnerApi.Coder originalMainInputCoder =
         stageCoderMap.get(inputPCollection.getPCollection().getCoderId());
-    String originalKeyCoderId = ModelCoders.getKvCoderComponents(originalCoder).keyCoderId();
-    assertThat(
-        stageCoderMap.get(originalKeyCoderId).getSpec().getUrn(),
-        is(CoderTranslation.JAVA_SERIALIZED_CODER_URN));
+    String originalKeyCoderId =
+        ModelCoders.getKvCoderComponents(originalMainInputCoder).keyCoderId();
+    RunnerApi.Coder originalKeyCoder = stageCoderMap.get(originalKeyCoderId);
+    assertThat(originalKeyCoder.getSpec().getUrn(), is(CoderTranslation.JAVA_SERIALIZED_CODER_URN));
 
     // Now create ProcessBundleDescriptor and check for the LengthPrefixCoder around the key coder
-    BeamFnApi.ProcessBundleDescriptor pbDescriptor =
+    BeamFnApi.ProcessBundleDescriptor pbd =
         ProcessBundleDescriptors.fromExecutableStage(
                 "test_stage", stage, Endpoints.ApiServiceDescriptor.getDefaultInstance())
             .getProcessBundleDescriptor();
+    Map<String, RunnerApi.Coder> pbsCoderMap = pbd.getCodersMap();
 
-    String inputPCollectionId = inputPCollection.getId();
-    String inputCoderId = pbDescriptor.getPcollectionsMap().get(inputPCollectionId).getCoderId();
+    RunnerApi.Coder pbsMainInputCoder =
+        pbsCoderMap.get(pbd.getPcollectionsOrThrow(inputPCollection.getId()).getCoderId());
+    String keyCoderId = ModelCoders.getKvCoderComponents(pbsMainInputCoder).keyCoderId();
+    RunnerApi.Coder keyCoder = pbsCoderMap.get(keyCoderId);
+    ensureLengthPrefixed(keyCoder, originalKeyCoder, pbsCoderMap);
 
-    Map<String, RunnerApi.Coder> pbCoderMap = pbDescriptor.getCodersMap();
-    RunnerApi.Coder coder = pbCoderMap.get(inputCoderId);
-    String keyCoderId = ModelCoders.getKvCoderComponents(coder).keyCoderId();
-
-    RunnerApi.Coder keyCoder = pbCoderMap.get(keyCoderId);
-    // Ensure length prefix
-    assertThat(keyCoder.getSpec().getUrn(), is(ModelCoders.LENGTH_PREFIX_CODER_URN));
-    String lengthPrefixWrappedCoderId = keyCoder.getComponentCoderIds(0);
+    TimerReference timerRef = Iterables.getOnlyElement(stage.getTimers());
+    String timerTransformId = timerRef.transform().getId();
+    RunnerApi.ParDoPayload parDoPayload =
+        RunnerApi.ParDoPayload.parseFrom(
+            pbd.getTransformsOrThrow(timerTransformId).getSpec().getPayload());
+    RunnerApi.TimerFamilySpec timerSpec =
+        parDoPayload.getTimerFamilySpecsOrThrow(timerRef.localName());
+    RunnerApi.Coder timerCoder = pbsCoderMap.get(timerSpec.getTimerFamilyCoderId());
+    String timerKeyCoderId = timerCoder.getComponentCoderIds(0);
+    RunnerApi.Coder timerKeyCoder = pbsCoderMap.get(timerKeyCoderId);
+    ensureLengthPrefixed(timerKeyCoder, originalKeyCoder, pbsCoderMap);
+  }
 
+  private static void ensureLengthPrefixed(
+      RunnerApi.Coder coder,
+      RunnerApi.Coder originalCoder,
+      Map<String, RunnerApi.Coder> pbsCoderMap) {
+    assertThat(coder.getSpec().getUrn(), is(ModelCoders.LENGTH_PREFIX_CODER_URN));
     // Check that the wrapped coder is unchanged
-    assertThat(lengthPrefixWrappedCoderId, is(originalKeyCoderId));
-    assertThat(
-        pbCoderMap.get(lengthPrefixWrappedCoderId), is(stageCoderMap.get(originalKeyCoderId)));
+    String lengthPrefixedWrappedCoderId = coder.getComponentCoderIds(0);
+    assertThat(pbsCoderMap.get(lengthPrefixedWrappedCoderId), is(originalCoder));
   }
 }
diff --git a/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner_test.py b/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner_test.py
@@ -33,7 +33,9 @@
 import unittest
 import uuid
 from builtins import range
+from typing import Any
 from typing import Dict
+from typing import Tuple
 
 # patches unittest.TestCase to be python3 compatible
 import hamcrest  # pylint: disable=ungrouped-imports
@@ -397,15 +399,24 @@ def process_clear_timer(self):
       assert_that(actual, equal_to(expected))
 
   def test_pardo_state_timers(self):
-    self._run_pardo_state_timers(False)
+    self._run_pardo_state_timers(windowed=False)
 
-  def test_windowed_pardo_state_timers(self):
-    self._run_pardo_state_timers(True)
+  def test_pardo_state_timers_non_standard_coder(self):
+    self._run_pardo_state_timers(windowed=False, key_type=Any)
 
-  def _run_pardo_state_timers(self, windowed):
+  def test_windowed_pardo_state_timers(self):
+    self._run_pardo_state_timers(windowed=True)
+
+  def _run_pardo_state_timers(self, windowed, key_type=None):
+    """
+    :param windowed: If True, uses an interval window, otherwise a global window
+    :param key_type: Allows to override the inferred key type. This is useful to
+    test the use of non-standard coders, e.g. Python's FastPrimitivesCoder.
+    """
     state_spec = userstate.BagStateSpec('state', beam.coders.StrUtf8Coder())
     timer_spec = userstate.TimerSpec('timer', userstate.TimeDomain.WATERMARK)
     elements = list('abcdefgh')
+    key = 'key'
     buffer_size = 3
 
     class BufferDoFn(beam.DoFn):
@@ -456,7 +467,8 @@ def is_buffered_correctly(actual):
           | beam.Map(lambda e: window.TimestampedValue(e, ord(e) % 2))
           | beam.WindowInto(
               window.FixedWindows(1) if windowed else window.GlobalWindows())
-          | beam.Map(lambda x: ('key', x))
+          | beam.Map(lambda x: (key, x)).with_output_types(
+              Tuple[key_type if key_type else type(key), Any])
           | beam.ParDo(BufferDoFn()))
 
       assert_that(actual, is_buffered_correctly)

Original file line number	Diff line number	Diff line change
`@@ -115,7 +115,11 @@ public abstract static class WindowedValueCoderComponents {`
`115`	`115`	`}`
`116`	`116`
`117`	`117`	`public static KvCoderComponents getKvCoderComponents(Coder coder) {`
`118`		`- checkArgument(KV_CODER_URN.equals(coder.getSpec().getUrn()));`
	`118`	`+ checkArgument(`
	`119`	`+ KV_CODER_URN.equals(coder.getSpec().getUrn()),`
	`120`	`+ "Provided coder %s is not of type %s",`
	`121`	`+ coder.getSpec().getUrn(),`
	`122`	`+ KV_CODER_URN);`
`119`	`123`	`return new AutoValue_ModelCoders_KvCoderComponents(`
`120`	`124`	`coder.getComponentCoderIds(0), coder.getComponentCoderIds(1));`
`121`	`125`	`}`