Merge pull request #34883: Refactor: separate PairWithRestriction from FnApiDoFnRunner

Author: kennknowles

sdks/java/harness/build.gradle

@@ -45,6 +45,7 @@ applyJavaNature(
     'AssignWindowsRunner': 'https://github.com/typetools/checker-framework/issues/3794',
     'WindowMergingFnRunner': 'https://github.com/typetools/checker-framework/issues/3794',
     'FnApiDoFnRunner': 'https://github.com/typetools/checker-framework/issues/5436',
+    'SplittablePairWithRestrictionDoFnRunner': 'https://github.com/typetools/checker-framework/issues/5436',
   ],
   automaticModuleName: 'org.apache.beam.fn.harness',
   testShadowJar: true,

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java

@@ -149,7 +149,6 @@ public Map<String, PTransformRunnerFactory> getPTransformRunnerFactories() {
       Factory factory = new Factory();
       return ImmutableMap.<String, PTransformRunnerFactory>builder()
           .put(PTransformTranslation.PAR_DO_TRANSFORM_URN, factory)
-          .put(PTransformTranslation.SPLITTABLE_PAIR_WITH_RESTRICTION_URN, factory)
           .put(PTransformTranslation.SPLITTABLE_SPLIT_AND_SIZE_RESTRICTIONS_URN, factory)
           .put(PTransformTranslation.SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN, factory)
           .put(
@@ -375,7 +374,6 @@ public final void addRunnerForPTransform(Context context) {
         case PTransformTranslation.PAR_DO_TRANSFORM_URN:
           mainOutputTag = (TupleTag) ParDoTranslation.getMainOutputTag(parDoPayload);
           break;
-        case PTransformTranslation.SPLITTABLE_PAIR_WITH_RESTRICTION_URN:
         case PTransformTranslation.SPLITTABLE_SPLIT_AND_SIZE_RESTRICTIONS_URN:
         case PTransformTranslation.SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN:
           mainOutputTag =
@@ -498,8 +496,6 @@ public final void addRunnerForPTransform(Context context) {
       case PTransformTranslation.SPLITTABLE_PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS_URN:
         addStartFunction.accept(this::startBundle);
         break;
-      case PTransformTranslation.SPLITTABLE_PAIR_WITH_RESTRICTION_URN:
-        // startBundle should not be invoked
       case PTransformTranslation.SPLITTABLE_SPLIT_AND_SIZE_RESTRICTIONS_URN:
         // startBundle should not be invoked
       case PTransformTranslation.SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN:
@@ -525,18 +521,6 @@ public final void addRunnerForPTransform(Context context) {
           this.processContext = new NonWindowObservingProcessBundleContext();
         }
         break;
-      case PTransformTranslation.SPLITTABLE_PAIR_WITH_RESTRICTION_URN:
-        if (doFnSignature.getInitialRestriction().observesWindow()
-            || (doFnSignature.getInitialWatermarkEstimatorState() != null
-                && doFnSignature.getInitialWatermarkEstimatorState().observesWindow())
-            || !sideInputMapping.isEmpty()) {
-          mainInputConsumer = this::processElementForWindowObservingPairWithRestriction;
-          this.processContext = new WindowObservingProcessBundleContext();
-        } else {
-          mainInputConsumer = this::processElementForPairWithRestriction;
-          this.processContext = new NonWindowObservingProcessBundleContext();
-        }
-        break;
       case PTransformTranslation.SPLITTABLE_SPLIT_AND_SIZE_RESTRICTIONS_URN:
         if ((doFnSignature.splitRestriction() != null
                 && doFnSignature.splitRestriction().observesWindow())
@@ -667,8 +651,6 @@ public void accept(WindowedValue input) throws Exception {
       case PTransformTranslation.SPLITTABLE_PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS_URN:
         addFinishFunction.accept(this::finishBundle);
         break;
-      case PTransformTranslation.SPLITTABLE_PAIR_WITH_RESTRICTION_URN:
-        // finishBundle should not be invoked
       case PTransformTranslation.SPLITTABLE_SPLIT_AND_SIZE_RESTRICTIONS_URN:
         // finishBundle should not be invoked
       case PTransformTranslation.SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN:
@@ -815,57 +797,6 @@ private void processElementForWindowObservingParDo(WindowedValue<InputT> elem) {
     }
   }
 
-  private void processElementForPairWithRestriction(WindowedValue<InputT> elem) {
-    currentElement = elem;
-    try {
-      currentRestriction = doFnInvoker.invokeGetInitialRestriction(processContext);
-      outputTo(
-          mainOutputConsumer,
-          (WindowedValue)
-              elem.withValue(
-                  KV.of(
-                      elem.getValue(),
-                      KV.of(
-                          currentRestriction,
-                          doFnInvoker.invokeGetInitialWatermarkEstimatorState(processContext)))));
-    } finally {
-      currentElement = null;
-      currentRestriction = null;
-    }
-
-    this.stateAccessor.finalizeState();
-  }
-
-  private void processElementForWindowObservingPairWithRestriction(WindowedValue<InputT> elem) {
-    currentElement = elem;
-    try {
-      Iterator<BoundedWindow> windowIterator =
-          (Iterator<BoundedWindow>) elem.getWindows().iterator();
-      while (windowIterator.hasNext()) {
-        currentWindow = windowIterator.next();
-        currentRestriction = doFnInvoker.invokeGetInitialRestriction(processContext);
-        outputTo(
-            mainOutputConsumer,
-            (WindowedValue)
-                WindowedValue.of(
-                    KV.of(
-                        elem.getValue(),
-                        KV.of(
-                            currentRestriction,
-                            doFnInvoker.invokeGetInitialWatermarkEstimatorState(processContext))),
-                    currentElement.getTimestamp(),
-                    currentWindow,
-                    currentElement.getPane()));
-      }
-    } finally {
-      currentElement = null;
-      currentWindow = null;
-      currentRestriction = null;
-    }
-
-    this.stateAccessor.finalizeState();
-  }
-
   private void processElementForSplitRestriction(
       WindowedValue<KV<InputT, KV<RestrictionT, WatermarkEstimatorStateT>>> elem) {
     currentElement = elem.withValue(elem.getValue().getKey());