Merge pull request #34919: Refactor: separate SplitAndSizeRestriction from FnApiDoFnRunner

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/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignature.java

@@ -53,6 +53,7 @@
 import org.apache.beam.sdk.values.TypeDescriptor;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Predicates;
 import org.checkerframework.checker.nullness.qual.Nullable;
+import org.checkerframework.dataflow.qual.Pure;
 
 /**
  * Describes the signature of a {@link DoFn}, in particular, which features it uses, which extra
@@ -67,91 +68,118 @@
 })
 public abstract class DoFnSignature {
   /** Class of the original {@link DoFn} from which this signature was produced. */
+  @Pure
   public abstract Class<? extends DoFn<?, ?>> fnClass();
 
   /** Whether this {@link DoFn} does a bounded amount of work per element. */
+  @Pure
   public abstract PCollection.IsBounded isBoundedPerElement();
 
   /** Details about this {@link DoFn}'s {@link DoFn.ProcessElement} method. */
+  @Pure
   public abstract ProcessElementMethod processElement();
 
   /** Details about the state cells that this {@link DoFn} declares. Immutable. */
+  @Pure
   public abstract Map<String, StateDeclaration> stateDeclarations();
 
   /** Details about this {@link DoFn}'s {@link DoFn.StartBundle} method. */
+  @Pure
   public abstract @Nullable BundleMethod startBundle();
 
   /** Details about this {@link DoFn}'s {@link DoFn.FinishBundle} method. */
+  @Pure
   public abstract @Nullable BundleMethod finishBundle();
 
   /** Details about this {@link DoFn}'s {@link DoFn.Setup} method. */
+  @Pure
   public abstract @Nullable LifecycleMethod setup();
 
   /** Details about this {@link DoFn}'s {@link DoFn.Teardown} method. */
+  @Pure
   public abstract @Nullable LifecycleMethod teardown();
 
   /** Details about this {@link DoFn}'s {@link DoFn.OnWindowExpiration} method. */
+  @Pure
   public abstract @Nullable OnWindowExpirationMethod onWindowExpiration();
 
   /** Timer declarations present on the {@link DoFn} class. Immutable. */
+  @Pure
   public abstract Map<String, TimerDeclaration> timerDeclarations();
 
   /** TimerMap declarations present on the {@link DoFn} class. Immutable. */
+  @Pure
   public abstract Map<String, TimerFamilyDeclaration> timerFamilyDeclarations();
 
   /** Field access declaration. */
+  @Pure
   public abstract @Nullable Map<String, FieldAccessDeclaration> fieldAccessDeclarations();
 
   /** Details about this {@link DoFn}'s {@link DoFn.GetInitialRestriction} method. */
+  @Pure
   public abstract @Nullable GetInitialRestrictionMethod getInitialRestriction();
 
   /** Details about this {@link DoFn}'s {@link DoFn.SplitRestriction} method. */
+  @Pure
   public abstract @Nullable SplitRestrictionMethod splitRestriction();
 
   /** Details about this {@link DoFn}'s {@link TruncateRestriction} method. */
+  @Pure
   public abstract @Nullable TruncateRestrictionMethod truncateRestriction();
 
   /** Details about this {@link DoFn}'s {@link DoFn.GetRestrictionCoder} method. */
+  @Pure
   public abstract @Nullable GetRestrictionCoderMethod getRestrictionCoder();
 
   /** Details about this {@link DoFn}'s {@link DoFn.GetWatermarkEstimatorStateCoder} method. */
+  @Pure
   public abstract @Nullable GetWatermarkEstimatorStateCoderMethod getWatermarkEstimatorStateCoder();
 
   /** Details about this {@link DoFn}'s {@link DoFn.GetInitialWatermarkEstimatorState} method. */
+  @Pure
   public abstract @Nullable GetInitialWatermarkEstimatorStateMethod
       getInitialWatermarkEstimatorState();
 
   /** Details about this {@link DoFn}'s {@link DoFn.NewWatermarkEstimator} method. */
+  @Pure
   public abstract @Nullable NewWatermarkEstimatorMethod newWatermarkEstimator();
 
   /** Details about this {@link DoFn}'s {@link DoFn.NewTracker} method. */
+  @Pure
   public abstract @Nullable NewTrackerMethod newTracker();
 
   /** Details about this {@link DoFn}'s {@link DoFn.GetSize} method. */
+  @Pure
   public abstract @Nullable GetSizeMethod getSize();
 
   /** Details about this {@link DoFn}'s {@link DoFn.OnTimer} methods. */
+  @Pure
   public abstract @Nullable Map<String, OnTimerMethod> onTimerMethods();
 
   /** Details about this {@link DoFn}'s {@link DoFn.OnTimerFamily} methods. */
+  @Pure
   public abstract @Nullable Map<String, OnTimerFamilyMethod> onTimerFamilyMethods();
 
   /** @deprecated use {@link #usesState()}, it's cleaner */
   @Deprecated
+  @Pure
   public boolean isStateful() {
     return stateDeclarations().size() > 0;
   }
 
   /** Whether the {@link DoFn} described by this signature uses state. */
+  @Pure
   public boolean usesState() {
     return stateDeclarations().size() > 0;
   }
 
   /** Whether the {@link DoFn} described by this signature uses timers. */
+  @Pure
   public boolean usesTimers() {
     return timerDeclarations().size() > 0 || timerFamilyDeclarations().size() > 0;
   }
 
+  @Pure
   static Builder builder() {
     return new AutoValue_DoFnSignature.Builder();
   }
@@ -208,12 +236,14 @@ abstract Builder setFieldAccessDeclarations(
 
     abstract Builder setOnTimerFamilyMethods(Map<String, OnTimerFamilyMethod> onTimerFamilyMethods);
 
+    @Pure
     abstract DoFnSignature build();
   }
 
   /** A method delegated to an annotated method of an underlying {@link DoFn}. */
   public interface DoFnMethod {
     /** The annotated method itself. */
+    @Pure
     Method targetMethod();
   }
 

sdks/java/harness/build.gradle

@@ -45,7 +45,6 @@ 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,