Refactor: separate SplittableSplitAndSizeRestriction

This is a further step in untangling FnApiDoFnRunner which
implements a variety of loosely-related transforms, which
mostly have in common that they call a user's DoFn but otherwise
are disjoint.

Author: kennknowles

sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFn.java

@@ -1211,11 +1211,11 @@ public interface MultiOutputReceiver {
    *
    * <ul>
    *   <li>The return type {@code WatermarkEstimatorStateT} defines the watermark state type used
-   *       within this splittable DoFn. All other methods that use a {@link
-   *       WatermarkEstimatorState @WatermarkEstimatorState} parameter must use the same type that
-   *       is used here. It is suggested to use as narrow of a return type definition as possible
-   *       (for example prefer to use a square type over a shape type as a square is a type of a
-   *       shape).
+   *       within this splittable DoFn. The return type is allowed to be nullable. All other methods
+   *       that use a {@link WatermarkEstimatorState @WatermarkEstimatorState} parameter must use
+   *       the same type that is used here. It is suggested to use as narrow of a return type
+   *       definition as possible (for example prefer to use a square type over a shape type as a
+   *       square is a type of a shape).
    *   <li>If one of its arguments is tagged with the {@link Element} annotation, then it will be
    *       passed the current element being processed; the argument must be of type {@code InputT}.
    *       Note that automatic conversion of {@link Row}s and {@link FieldAccess} parameters are

sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnInvoker.java

@@ -43,6 +43,7 @@
 import org.apache.beam.sdk.transforms.windowing.PaneInfo;
 import org.apache.beam.sdk.values.Row;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.Instant;
 
 /**
@@ -96,7 +97,7 @@ <RestrictionT> RestrictionT invokeGetInitialRestriction(
   void invokeSplitRestriction(ArgumentProvider<InputT, OutputT> arguments);
 
   /** Invoke the {@link TruncateRestriction} method on the bound {@link DoFn}. */
-  <RestrictionT> TruncateResult<RestrictionT> invokeTruncateRestriction(
+  <RestrictionT> @Nullable TruncateResult<RestrictionT> invokeTruncateRestriction(
       ArgumentProvider<InputT, OutputT> arguments);
 
   /**
@@ -117,7 +118,13 @@ <RestrictionT> TruncateResult<RestrictionT> invokeTruncateRestriction(
   <RestrictionT, PositionT> RestrictionTracker<RestrictionT, PositionT> invokeNewTracker(
       ArgumentProvider<InputT, OutputT> arguments);
 
-  /** Invoke the {@link DoFn.NewWatermarkEstimator} method on the bound {@link DoFn}. */
+  /**
+   * Invoke the {@link DoFn.NewWatermarkEstimator} method on the bound {@link DoFn}.
+   *
+   * <p>Note that since {@code WatermarkEstimatorStateT} is permitted to be a nullable type, if this
+   * method returns {@code null} that is interpreted as a valid watermark estimator state, not the
+   * absence of a state.
+   */
   @SuppressWarnings("TypeParameterUnusedInFormals")
   <WatermarkEstimatorStateT>
       WatermarkEstimator<WatermarkEstimatorStateT> invokeNewWatermarkEstimator(

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

@@ -22,7 +22,6 @@
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
 
 import com.google.auto.service.AutoService;
-import com.google.auto.value.AutoValue;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -57,6 +56,7 @@
 import org.apache.beam.runners.core.metrics.ShortIdMap;
 import org.apache.beam.runners.core.metrics.SimpleMonitoringInfoBuilder;
 import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.annotations.Internal;
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.coders.DoubleCoder;
 import org.apache.beam.sdk.coders.IterableCoder;
@@ -134,6 +134,7 @@
 @SuppressWarnings({
   "rawtypes" // TODO(https://github.com/apache/beam/issues/20447)
 })
+@Internal
 public class FnApiDoFnRunner<InputT, RestrictionT, PositionT, WatermarkEstimatorStateT, OutputT>
     implements FnApiStateAccessor.MutatingStateContext<Object, BoundedWindow> {
   /** A registrar which provides a factory to handle Java {@link DoFn}s. */
@@ -144,7 +145,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_SPLIT_AND_SIZE_RESTRICTIONS_URN, factory)
           .put(PTransformTranslation.SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN, factory)
           .put(
               PTransformTranslation.SPLITTABLE_PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS_URN, factory)
@@ -158,6 +158,9 @@ static class Factory<InputT, RestrictionT, PositionT, WatermarkEstimatorStateT,
     @Override
     public final void addRunnerForPTransform(Context context) throws IOException {
 
+      FnApiStateAccessor<Object> stateAccessor =
+          FnApiStateAccessor.Factory.factoryForPTransformContext(context).create();
+
       FnApiDoFnRunner<InputT, RestrictionT, PositionT, WatermarkEstimatorStateT, OutputT> runner =
           new FnApiDoFnRunner<>(
               context.getPipelineOptions(),
@@ -175,7 +178,9 @@ public final void addRunnerForPTransform(Context context) throws IOException {
               context::addBundleProgressReporter,
               context.getSplitListener(),
               context.getBundleFinalizer(),
-              FnApiStateAccessor.Factory.factoryForPTransformContext(context));
+              stateAccessor);
+
+      stateAccessor.setKeyAndWindowContext(runner);
 
       for (Map.Entry<String, KV<TimeDomain, Coder<Timer<Object>>>> entry :
           runner.timerFamilyInfos.entrySet()) {
@@ -338,7 +343,7 @@ public final void addRunnerForPTransform(Context context) throws IOException {
       Consumer<BundleProgressReporter> addBundleProgressReporter,
       BundleSplitListener splitListener,
       BundleFinalizer bundleFinalizer,
-      FnApiStateAccessor.Factory<Object> stateAccessorFactory) {
+      FnApiStateAccessor<Object> stateAccessor) {
     this.pipelineOptions = pipelineOptions;
     this.pTransformId = pTransformId;
     this.pTransform = pTransform;
@@ -480,24 +485,6 @@ public final void addRunnerForPTransform(Context context) throws IOException {
           this.processContext = new NonWindowObservingProcessBundleContext();
         }
         break;
-      case PTransformTranslation.SPLITTABLE_SPLIT_AND_SIZE_RESTRICTIONS_URN:
-        if ((doFnSignature.splitRestriction() != null
-                && doFnSignature.splitRestriction().observesWindow())
-            || (doFnSignature.newTracker() != null && doFnSignature.newTracker().observesWindow())
-            || (doFnSignature.getSize() != null && doFnSignature.getSize().observesWindow())
-            || !sideInputMapping.isEmpty()) {
-          mainInputConsumer = this::processElementForWindowObservingSplitRestriction;
-          this.processContext =
-              new SizedRestrictionWindowObservingProcessBundleContext(
-                  PTransformTranslation.SPLITTABLE_SPLIT_AND_SIZE_RESTRICTIONS_URN);
-
-        } else {
-          mainInputConsumer = this::processElementForSplitRestriction;
-          this.processContext =
-              new SizedRestrictionNonWindowObservingProcessBundleContext(
-                  PTransformTranslation.SPLITTABLE_SPLIT_AND_SIZE_RESTRICTIONS_URN);
-        }
-        break;
       case PTransformTranslation.SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN:
         if ((doFnSignature.truncateRestriction() != null
                 && doFnSignature.truncateRestriction().observesWindow())
@@ -675,7 +662,7 @@ private ByteString encodeProgress(double value) throws IOException {
         // no-op
     }
 
-    this.stateAccessor = stateAccessorFactory.create(this);
+    this.stateAccessor = stateAccessor;
 
     // Register as a consumer for each timer.
     this.outboundTimerReceivers = new HashMap<>();
@@ -748,80 +735,6 @@ private void processElementForWindowObservingParDo(WindowedValue<InputT> elem) {
     }
   }
 
-  private void processElementForSplitRestriction(
-      WindowedValue<KV<InputT, KV<RestrictionT, WatermarkEstimatorStateT>>> elem) {
-    currentElement = elem.withValue(elem.getValue().getKey());
-    currentRestriction = elem.getValue().getValue().getKey();
-    currentWatermarkEstimatorState = elem.getValue().getValue().getValue();
-    currentTrackerClaimed = new AtomicBoolean(false);
-    currentTracker =
-        RestrictionTrackers.observe(
-            doFnInvoker.invokeNewTracker(processContext),
-            new ClaimObserver<PositionT>() {
-              private final AtomicBoolean claimed =
-                  Preconditions.checkNotNull(currentTrackerClaimed);
-
-              @Override
-              public void onClaimed(PositionT position) {
-                claimed.lazySet(true);
-              }
-
-              @Override
-              public void onClaimFailed(PositionT position) {}
-            });
-    try {
-      doFnInvoker.invokeSplitRestriction(processContext);
-    } finally {
-      currentElement = null;
-      currentRestriction = null;
-      currentWatermarkEstimatorState = null;
-      currentTracker = null;
-      currentTrackerClaimed = null;
-    }
-
-    this.stateAccessor.finalizeState();
-  }
-
-  private void processElementForWindowObservingSplitRestriction(
-      WindowedValue<KV<InputT, KV<RestrictionT, WatermarkEstimatorStateT>>> elem) {
-    currentElement = elem.withValue(elem.getValue().getKey());
-    currentRestriction = elem.getValue().getValue().getKey();
-    currentWatermarkEstimatorState = elem.getValue().getValue().getValue();
-    try {
-      Iterator<BoundedWindow> windowIterator =
-          (Iterator<BoundedWindow>) elem.getWindows().iterator();
-      while (windowIterator.hasNext()) {
-        currentWindow = windowIterator.next();
-        currentTrackerClaimed = new AtomicBoolean(false);
-        currentTracker =
-            RestrictionTrackers.observe(
-                doFnInvoker.invokeNewTracker(processContext),
-                new ClaimObserver<PositionT>() {
-                  private final AtomicBoolean claimed =
-                      Preconditions.checkNotNull(currentTrackerClaimed);
-
-                  @Override
-                  public void onClaimed(PositionT position) {
-                    claimed.lazySet(true);
-                  }
-
-                  @Override
-                  public void onClaimFailed(PositionT position) {}
-                });
-        doFnInvoker.invokeSplitRestriction(processContext);
-      }
-    } finally {
-      currentElement = null;
-      currentRestriction = null;
-      currentWatermarkEstimatorState = null;
-      currentWindow = null;
-      currentTracker = null;
-      currentTrackerClaimed = null;
-    }
-
-    this.stateAccessor.finalizeState();
-  }
-
   private void processElementForTruncateRestriction(
       WindowedValue<KV<KV<InputT, KV<RestrictionT, WatermarkEstimatorStateT>>, Double>> elem) {
     currentElement = elem.withValue(elem.getValue().getKey().getKey());
@@ -906,49 +819,6 @@ public void onClaimFailed(PositionT position) {}
     this.stateAccessor.finalizeState();
   }
 
-  /** Internal class to hold the primary and residual roots when converted to an input element. */
-  @AutoValue
-  @AutoValue.CopyAnnotations
-  abstract static class WindowedSplitResult {
-    public static WindowedSplitResult forRoots(
-        WindowedValue<?> primaryInFullyProcessedWindowsRoot,
-        WindowedValue<?> primarySplitRoot,
-        WindowedValue<?> residualSplitRoot,
-        WindowedValue<?> residualInUnprocessedWindowsRoot) {
-      return new AutoValue_FnApiDoFnRunner_WindowedSplitResult(
-          primaryInFullyProcessedWindowsRoot,
-          primarySplitRoot,
-          residualSplitRoot,
-          residualInUnprocessedWindowsRoot);
-    }
-
-    public abstract @Nullable WindowedValue<?> getPrimaryInFullyProcessedWindowsRoot();
-
-    public abstract @Nullable WindowedValue<?> getPrimarySplitRoot();
-
-    public abstract @Nullable WindowedValue<?> getResidualSplitRoot();
-
-    public abstract @Nullable WindowedValue<?> getResidualInUnprocessedWindowsRoot();
-  }
-
-  @AutoValue
-  @AutoValue.CopyAnnotations
-  abstract static class SplitResultsWithStopIndex {
-    public static SplitResultsWithStopIndex of(
-        WindowedSplitResult windowSplit,
-        HandlesSplits.SplitResult downstreamSplit,
-        int newWindowStopIndex) {
-      return new AutoValue_FnApiDoFnRunner_SplitResultsWithStopIndex(
-          windowSplit, downstreamSplit, newWindowStopIndex);
-    }
-
-    public abstract @Nullable WindowedSplitResult getWindowSplit();
-
-    public abstract HandlesSplits.@Nullable SplitResult getDownstreamSplit();
-
-    public abstract int getNewWindowStopIndex();
-  }
-
   private void processElementForWindowObservingSizedElementAndRestriction(
       WindowedValue<KV<KV<InputT, KV<RestrictionT, WatermarkEstimatorStateT>>, Double>> elem) {
     currentElement = elem.withValue(elem.getValue().getKey().getKey());

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

@@ -0,0 +1,41 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.fn.harness;
+
+import com.google.auto.value.AutoValue;
+import org.apache.beam.sdk.annotations.Internal;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+@AutoValue
+@AutoValue.CopyAnnotations
+@Internal
+abstract class SplitResultsWithStopIndex {
+  public static SplitResultsWithStopIndex of(
+      WindowedSplitResult windowSplit,
+      HandlesSplits.SplitResult downstreamSplit,
+      int newWindowStopIndex) {
+    return new AutoValue_SplitResultsWithStopIndex(
+        windowSplit, downstreamSplit, newWindowStopIndex);
+  }
+
+  public abstract @Nullable WindowedSplitResult getWindowSplit();
+
+  public abstract HandlesSplits.@Nullable SplitResult getDownstreamSplit();
+
+  public abstract int getNewWindowStopIndex();
+}