[#20502] Remove nullness suppressions from the direct-java runner

runners/direct-java/src/main/java/org/apache/beam/runners/direct/BoundedReadEvaluatorFactory.java

@@ -17,6 +17,8 @@
  */
 package org.apache.beam.runners.direct;
 
+import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
+
 import com.google.auto.value.AutoValue;
 import java.io.IOException;
 import java.util.Collection;
@@ -52,8 +54,7 @@
  * for the {@link PrimitiveBoundedRead SplittableParDo.PrimitiveBoundedRead} {@link PTransform}.
  */
 @SuppressWarnings({
-  "rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
-  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
+  "rawtypes" // TODO(https://github.com/apache/beam/issues/20447)
 })
 final class BoundedReadEvaluatorFactory implements TransformEvaluatorFactory {
   /**
@@ -149,7 +150,8 @@ public void processElement(WindowedValue<BoundedSourceShard<OutputT>> element)
       BoundedSource<OutputT> source = element.getValue().getSource();
       try (final BoundedReader<OutputT> reader = source.createReader(options)) {
         boolean contentsRemaining = reader.start();
-        Future<BoundedSource<OutputT>> residualFuture = startDynamicSplitThread(source, reader);
+        Future<@Nullable BoundedSource<OutputT>> residualFuture =
+            startDynamicSplitThread(source, reader);
         UncommittedBundle<OutputT> output = evaluationContext.createBundle(outputPCollection);
         while (contentsRemaining) {
           output.add(
@@ -159,24 +161,24 @@ public void processElement(WindowedValue<BoundedSourceShard<OutputT>> element)
         }
         resultBuilder.addOutput(output);
         try {
-          BoundedSource<OutputT> residual = residualFuture.get();
+          @Nullable BoundedSource<OutputT> residual = residualFuture.get();
           if (residual != null) {
             resultBuilder.addUnprocessedElements(
                 element.withValue(BoundedSourceShard.of(residual)));
           }
         } catch (ExecutionException exex) {
           // Un-and-rewrap the exception thrown by attempting to split
-          throw UserCodeException.wrap(exex.getCause());
+          throw UserCodeException.wrap(checkStateNotNull(exex.getCause()));
         }
       }
     }
 
-    private Future<BoundedSource<OutputT>> startDynamicSplitThread(
+    private Future<@Nullable BoundedSource<OutputT>> startDynamicSplitThread(
         BoundedSource<OutputT> source, BoundedReader<OutputT> reader) throws Exception {
       if (source.getEstimatedSizeBytes(options) > minimumDynamicSplitSize) {
         return produceSplitExecutor.submit(new GenerateSplitAtHalfwayPoint<>(reader));
       } else {
-        SettableFuture<BoundedSource<OutputT>> emptyFuture = SettableFuture.create();
+        SettableFuture<@Nullable BoundedSource<OutputT>> emptyFuture = SettableFuture.create();
         emptyFuture.set(null);
         return emptyFuture;
       }
@@ -230,15 +232,16 @@ public Collection<CommittedBundle<BoundedSourceShard<T>>> getInitialInputs(
     }
   }
 
-  private static class GenerateSplitAtHalfwayPoint<T> implements Callable<BoundedSource<T>> {
+  private static class GenerateSplitAtHalfwayPoint<T>
+      implements Callable<@Nullable BoundedSource<T>> {
     private final BoundedReader<T> reader;
 
     private GenerateSplitAtHalfwayPoint(BoundedReader<T> reader) {
       this.reader = reader;
     }
 
     @Override
-    public BoundedSource<T> call() throws Exception {
+    public @Nullable BoundedSource<T> call() throws Exception {
       // Splits at halfway of the remaining work.
       Double currentlyConsumed = reader.getFractionConsumed();
       if (currentlyConsumed == null || currentlyConsumed == 1.0) {

runners/direct-java/src/main/java/org/apache/beam/runners/direct/CloningBundleFactory.java

@@ -17,6 +17,8 @@
  */
 package org.apache.beam.runners.direct;
 
+import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
+
 import org.apache.beam.runners.local.StructuralKey;
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.coders.CoderException;
@@ -25,15 +27,13 @@
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.sdk.values.WindowedValue;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects;
+import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.Instant;
 
 /**
  * A {@link BundleFactory} where a created {@link UncommittedBundle} clones all elements added to it
  * using the coder of the {@link PCollection}.
  */
-@SuppressWarnings({
-  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
 class CloningBundleFactory implements BundleFactory {
   private static final CloningBundleFactory INSTANCE = new CloningBundleFactory();
 
@@ -70,11 +70,12 @@ private static class CloningBundle<T> implements UncommittedBundle<T> {
 
     private CloningBundle(UncommittedBundle<T> underlying) {
       this.underlying = underlying;
-      this.coder = underlying.getPCollection().getCoder();
+      // CloningBundle only wraps non-root bundles, whose PCollection is always present.
+      this.coder = checkStateNotNull(underlying.getPCollection()).getCoder();
     }
 
     @Override
-    public PCollection<T> getPCollection() {
+    public @Nullable PCollection<T> getPCollection() {
       return underlying.getPCollection();
     }
 

runners/direct-java/src/main/java/org/apache/beam/runners/direct/CommittedResult.java

@@ -24,9 +24,6 @@
 
 /** A {@link TransformResult} that has been committed. */
 @AutoValue
-@SuppressWarnings({
-  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
 abstract class CommittedResult<ExecutableT> {
   /** Returns the {@link AppliedPTransform} that produced this result. */
   public abstract ExecutableT getExecutable();

runners/direct-java/src/main/java/org/apache/beam/runners/direct/CopyOnAccessInMemoryStateInternals.java

@@ -56,6 +56,7 @@
 import org.apache.beam.sdk.transforms.windowing.TimestampCombiner;
 import org.apache.beam.sdk.util.CombineFnUtil;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
+import org.checkerframework.checker.nullness.qual.NonNull;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.Instant;
 
@@ -65,10 +66,9 @@
  * accessed, an independent copy will be created within this table.
  */
 @SuppressWarnings({
-  "rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
-  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
+  "rawtypes" // TODO(https://github.com/apache/beam/issues/20447)
 })
-class CopyOnAccessInMemoryStateInternals<K> implements StateInternals {
+class CopyOnAccessInMemoryStateInternals<K extends @NonNull Object> implements StateInternals {
   private final CopyOnAccessInMemoryStateTable table;
 
   private K key;
@@ -77,12 +77,13 @@ class CopyOnAccessInMemoryStateInternals<K> implements StateInternals {
    * Creates a new {@link CopyOnAccessInMemoryStateInternals} with the underlying (possibly null)
    * StateInternals.
    */
-  public static <K> CopyOnAccessInMemoryStateInternals withUnderlying(
+  public static <K extends @NonNull Object> CopyOnAccessInMemoryStateInternals withUnderlying(
       K key, @Nullable CopyOnAccessInMemoryStateInternals underlying) {
     return new CopyOnAccessInMemoryStateInternals<>(key, underlying);
   }
 
-  private CopyOnAccessInMemoryStateInternals(K key, CopyOnAccessInMemoryStateInternals underlying) {
+  private CopyOnAccessInMemoryStateInternals(
+      K key, @Nullable CopyOnAccessInMemoryStateInternals underlying) {
     this.key = key;
     table = new CopyOnAccessInMemoryStateTable(underlying == null ? null : underlying.table);
   }
@@ -169,7 +170,7 @@ private static class CopyOnAccessInMemoryStateTable extends StateTable {
     /** The earliest watermark hold in this table. */
     private Optional<Instant> earliestWatermarkHold;
 
-    public CopyOnAccessInMemoryStateTable(StateTable underlying) {
+    public CopyOnAccessInMemoryStateTable(@Nullable StateTable underlying) {
       this.underlying = Optional.ofNullable(underlying);
       binderFactory = new CopyOnBindBinderFactory(this.underlying);
       earliestWatermarkHold = Optional.empty();

runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectExecutionContext.java

@@ -28,6 +28,7 @@
 import org.apache.beam.runners.direct.WatermarkManager.TransformWatermarks;
 import org.apache.beam.runners.local.StructuralKey;
 import org.apache.beam.sdk.transforms.DoFn.BundleFinalizer;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 /**
  * Execution Context for the {@link DirectRunner}.
@@ -36,21 +37,20 @@
  * for each thread that requires it.
  */
 @SuppressWarnings({
-  "rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
-  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
+  "rawtypes" // TODO(https://github.com/apache/beam/issues/20447)
 })
 class DirectExecutionContext {
 
   private final Clock clock;
   private final StructuralKey<?> key;
-  private final CopyOnAccessInMemoryStateInternals existingState;
+  private final @Nullable CopyOnAccessInMemoryStateInternals existingState;
   private final TransformWatermarks watermarks;
   private Map<String, DirectStepContext> cachedStepContexts = new LinkedHashMap<>();
 
   public DirectExecutionContext(
       Clock clock,
       StructuralKey<?> key,
-      CopyOnAccessInMemoryStateInternals existingState,
+      @Nullable CopyOnAccessInMemoryStateInternals existingState,
       TransformWatermarks watermarks) {
     this.clock = clock;
     this.key = key;
@@ -69,9 +69,9 @@ public DirectStepContext getStepContext(String stepName) {
 
   /** Step Context for the {@link DirectRunner}. */
   public class DirectStepContext implements StepContext {
-    private CopyOnAccessInMemoryStateInternals<?> stateInternals;
-    private DirectTimerInternals timerInternals;
-    private InMemoryBundleFinalizer bundleFinalizer;
+    private @Nullable CopyOnAccessInMemoryStateInternals<?> stateInternals;
+    private @Nullable DirectTimerInternals timerInternals;
+    private @Nullable InMemoryBundleFinalizer bundleFinalizer;
 
     public DirectStepContext() {}
 
@@ -111,7 +111,7 @@ public List<InMemoryBundleFinalizer.Finalization> getAndClearFinalizations() {
      * Commits the state of this step, and returns the committed state. If the step has not accessed
      * any state, return null.
      */
-    public CopyOnAccessInMemoryStateInternals commitState() {
+    public @Nullable CopyOnAccessInMemoryStateInternals commitState() {
       if (stateInternals != null) {
         return stateInternals.commit();
       }

runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGraph.java

@@ -17,6 +17,8 @@
  */
 package org.apache.beam.runners.direct;
 
+import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
+
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
@@ -34,7 +36,7 @@
  * Methods for interacting with the underlying structure of a {@link Pipeline} that is being
  * executed with the {@link DirectRunner}.
  */
-@SuppressWarnings({"keyfor", "nullness"}) // TODO(https://github.com/apache/beam/issues/20497)
+@SuppressWarnings({"keyfor"}) // TODO(https://github.com/apache/beam/issues/20497)
 class DirectGraph implements ExecutableGraph<AppliedPTransform<?, ?, ?>, PValue> {
   private final Map<PCollection<?>, AppliedPTransform<?, ?, ?>> producers;
   private final Map<PCollectionView<?>, AppliedPTransform<?, ?, ?>> viewWriters;
@@ -68,7 +70,7 @@ private DirectGraph(
   @Override
   public AppliedPTransform<?, ?, ?> getProducer(PValue produced) {
     if (produced instanceof PCollection) {
-      return producers.get(produced);
+      return checkStateNotNull(producers.get(produced), "No producer for %s", produced);
     } else if (produced instanceof PCollectionView) {
       return getWriter((PCollectionView<?>) produced);
     }
@@ -96,7 +98,7 @@ public Collection<PValue> getPerElementInputs(AppliedPTransform<?, ?, ?> transfo
   }
 
   private AppliedPTransform<?, ?, ?> getWriter(PCollectionView<?> view) {
-    return viewWriters.get(view);
+    return checkStateNotNull(viewWriters.get(view), "No writer for %s", view);
   }
 
   @Override
@@ -119,6 +121,6 @@ Set<PCollectionView<?>> getViews() {
   }
 
   String getStepName(AppliedPTransform<?, ?, ?> step) {
-    return stepNames.get(step);
+    return checkStateNotNull(stepNames.get(step), "No step name for %s", step);
   }
 }

runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGraphVisitor.java

@@ -47,9 +47,6 @@
  * {@link Pipeline}. This is used to schedule consuming {@link PTransform PTransforms} to consume
  * input after the upstream transform has produced and committed output.
  */
-@SuppressWarnings({
-  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
 class DirectGraphVisitor extends PipelineVisitor.Defaults {
   private static final Logger LOG = LoggerFactory.getLogger(DirectGraphVisitor.class);
 
@@ -119,12 +116,11 @@ public void visitPrimitiveTransform(TransformHierarchy.Node node) {
         allConsumers.put(value, appliedTransform);
       }
     }
-    if (node.getTransform() instanceof ParDo.MultiOutput) {
-      consumedViews.addAll(
-          ((ParDo.MultiOutput<?, ?>) node.getTransform()).getSideInputs().values());
-    } else if (node.getTransform() instanceof WriteView) {
-      viewWriters.put(
-          ((WriteView) node.getTransform()).getView(), node.toAppliedPTransform(getPipeline()));
+    PTransform<?, ?> transform = node.getTransform();
+    if (transform instanceof ParDo.MultiOutput) {
+      consumedViews.addAll(((ParDo.MultiOutput<?, ?>) transform).getSideInputs().values());
+    } else if (transform instanceof WriteView) {
+      viewWriters.put(((WriteView) transform).getView(), node.toAppliedPTransform(getPipeline()));
     }
   }
 

runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectMetrics.java

@@ -47,12 +47,10 @@
 import org.apache.beam.sdk.metrics.MetricsFilter;
 import org.apache.beam.sdk.metrics.StringSetResult;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+import org.checkerframework.checker.nullness.qual.NonNull;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 /** Implementation of {@link MetricResults} for the Direct Runner. */
-@SuppressWarnings({
-  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
 class DirectMetrics extends MetricResults {
 
   private interface MetricAggregation<UpdateT, ResultT> {
@@ -69,7 +67,7 @@ private interface MetricAggregation<UpdateT, ResultT> {
    * @param <UpdateT> The type of raw data received and aggregated across updates.
    * @param <ResultT> The type of result extracted from the data.
    */
-  private static class DirectMetric<UpdateT, ResultT> {
+  private static class DirectMetric<UpdateT extends @NonNull Object, ResultT> {
     private final MetricAggregation<UpdateT, ResultT> aggregation;
 
     private final Executor executor;
@@ -326,10 +324,10 @@ public MetricQueryResults queryMetrics(@Nullable MetricsFilter filter) {
   }
 
   private <ResultT> void maybeExtractResult(
-      MetricsFilter filter,
+      @Nullable MetricsFilter filter,
       ImmutableList.Builder<MetricResult<ResultT>> resultsBuilder,
       Map.Entry<MetricKey, ? extends DirectMetric<?, ResultT>> entry) {
-    if (MetricFiltering.matches(filter, entry.getKey())) {
+    if (filter == null || MetricFiltering.matches(filter, entry.getKey())) {
       resultsBuilder.add(
           MetricResult.create(
               entry.getKey(),

runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectRunner.java

@@ -17,6 +17,8 @@
  */
 package org.apache.beam.runners.direct;
 
+import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
+
 import com.fasterxml.jackson.databind.ObjectMapper;
 import java.io.IOException;
 import java.util.Collection;
@@ -66,9 +68,6 @@
  * contained within a {@link Pipeline} does not break assumptions within the Beam model, to improve
  * the ability to execute a {@link Pipeline} at scale on a distributed backend.
  */
-@SuppressWarnings({
-  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
 public class DirectRunner extends PipelineRunner<DirectPipelineResult> {
 
   enum Enforcement {
@@ -222,7 +221,7 @@ public DirectPipelineResult run(Pipeline pipeline) {
         try {
           result.waitUntilFinish();
         } catch (UserCodeException userException) {
-          throw new PipelineExecutionException(userException.getCause());
+          throw new PipelineExecutionException(checkStateNotNull(userException.getCause()));
         } catch (RuntimeException | OutOfMemoryError e) {
           throw e;
         } catch (Throwable t) {
@@ -368,6 +367,10 @@ public State cancel() {
      * org.apache.beam.sdk.PipelineResult.State#FAILED}.
      */
     @Override
+    // PipelineResult.waitUntilFinish(Duration) is documented to return null on timeout, but its
+    // interface return type is not annotated @Nullable; annotating it would ripple across every
+    // runner's override, so the timeout null is suppressed here rather than in the sdk-core API.
+    @SuppressWarnings("nullness")
     public State waitUntilFinish(Duration duration) {
       if (this.state.isTerminal()) {
         return this.state;
@@ -379,7 +382,7 @@ public State waitUntilFinish(Duration duration) {
         // Emulates the behavior of Pipeline#run(), where a stack trace caused by a
         // UserCodeException is truncated and replaced with the stack starting at the call to
         // waitToFinish
-        throw new Pipeline.PipelineExecutionException(uce.getCause());
+        throw new Pipeline.PipelineExecutionException(checkStateNotNull(uce.getCause()));
       } catch (Exception e) {
         if (e instanceof InterruptedException) {
           Thread.currentThread().interrupt();