Introduce ValueWithMetadata and its receivers

When we introduce OutputBuilder, it will implement ValueWithMetadata,
and every time we create an OutputBuilder we will want to provide

1. the default metadata for the output
2. the receiver for the output

This change unblocks 2 and reduces the size of the final change.

Author: kennknowles

runners/core-java/src/main/java/org/apache/beam/runners/core/DoFnRunners.java

@@ -29,6 +29,7 @@
 import org.apache.beam.sdk.transforms.DoFnSchemaInformation;
 import org.apache.beam.sdk.transforms.reflect.DoFnSignatures;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.util.ValueWithMetadataMultiReceiver;
 import org.apache.beam.sdk.util.WindowedValue;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.PCollectionView;
@@ -41,12 +42,6 @@
   "nullness" // TODO(https://github.com/apache/beam/issues/20497)
 })
 public class DoFnRunners {
-  /** Information about how to create output receivers and output to them. */
-  public interface OutputManager {
-    /** Outputs a single element to the receiver indicated by the given {@link TupleTag}. */
-    <T> void output(TupleTag<T> tag, WindowedValue<T> output);
-  }
-
   /**
    * Returns an implementation of {@link DoFnRunner} that for a {@link DoFn}.
    *
@@ -58,7 +53,7 @@ public static <InputT, OutputT> DoFnRunner<InputT, OutputT> simpleRunner(
       PipelineOptions options,
       DoFn<InputT, OutputT> fn,
       SideInputReader sideInputReader,
-      OutputManager outputManager,
+      ValueWithMetadataMultiReceiver outputManager,
       TupleTag<OutputT> mainOutputTag,
       List<TupleTag<?>> additionalOutputTags,
       StepContext stepContext,
@@ -168,7 +163,7 @@ ProcessFnRunner<InputT, OutputT, RestrictionT> newProcessFnRunner(
           PipelineOptions options,
           Collection<PCollectionView<?>> views,
           ReadyCheckingSideInputReader sideInputReader,
-          OutputManager outputManager,
+          ValueWithMetadataMultiReceiver outputManager,
           TupleTag<OutputT> mainOutputTag,
           List<TupleTag<?>> additionalOutputTags,
           StepContext stepContext,

runners/core-java/src/main/java/org/apache/beam/runners/core/GroupAlsoByWindowViaWindowSetNewDoFn.java

@@ -17,20 +17,17 @@
  */
 package org.apache.beam.runners.core;
 
-import java.util.Collection;
 import org.apache.beam.model.pipeline.v1.RunnerApi;
 import org.apache.beam.runners.core.triggers.ExecutableTriggerStateMachine;
 import org.apache.beam.runners.core.triggers.TriggerStateMachines;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
-import org.apache.beam.sdk.transforms.windowing.PaneInfo;
 import org.apache.beam.sdk.util.SystemDoFnInternal;
-import org.apache.beam.sdk.util.WindowedValue;
+import org.apache.beam.sdk.util.ValueWithMetadataMultiReceiver;
 import org.apache.beam.sdk.util.construction.TriggerTranslation;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.TupleTag;
 import org.apache.beam.sdk.values.WindowingStrategy;
-import org.joda.time.Instant;
 
 /**
  * A general {@link GroupAlsoByWindowsAggregators}. This delegates all of the logic to the {@link
@@ -51,7 +48,7 @@ DoFn<KeyedWorkItem<K, InputT>, KV<K, OutputT>> create(
           TimerInternalsFactory<K> timerInternalsFactory,
           SideInputReader sideInputReader,
           SystemReduceFn<K, InputT, ?, OutputT, W> reduceFn,
-          DoFnRunners.OutputManager outputManager,
+          ValueWithMetadataMultiReceiver outputManager,
           TupleTag<KV<K, OutputT>> mainTag) {
     return new GroupAlsoByWindowViaWindowSetNewDoFn<>(
         strategy,
@@ -68,7 +65,7 @@ DoFn<KeyedWorkItem<K, InputT>, KV<K, OutputT>> create(
   private transient StateInternalsFactory<K> stateInternalsFactory;
   private transient TimerInternalsFactory<K> timerInternalsFactory;
   private transient SideInputReader sideInputReader;
-  private transient DoFnRunners.OutputManager outputManager;
+  private transient ValueWithMetadataMultiReceiver outputManager;
   private TupleTag<KV<K, OutputT>> mainTag;
 
   public GroupAlsoByWindowViaWindowSetNewDoFn(
@@ -77,7 +74,7 @@ public GroupAlsoByWindowViaWindowSetNewDoFn(
       TimerInternalsFactory<K> timerInternalsFactory,
       SideInputReader sideInputReader,
       SystemReduceFn<K, InputT, ?, OutputT, W> reduceFn,
-      DoFnRunners.OutputManager outputManager,
+      ValueWithMetadataMultiReceiver outputManager,
       TupleTag<KV<K, OutputT>> mainTag) {
     this.timerInternalsFactory = timerInternalsFactory;
     this.sideInputReader = sideInputReader;
@@ -91,29 +88,6 @@ public GroupAlsoByWindowViaWindowSetNewDoFn(
     this.triggerProto = TriggerTranslation.toProto(windowingStrategy.getTrigger());
   }
 
-  private OutputWindowedValue<KV<K, OutputT>> outputWindowedValue() {
-    return new OutputWindowedValue<KV<K, OutputT>>() {
-      @Override
-      public void outputWindowedValue(
-          KV<K, OutputT> output,
-          Instant timestamp,
-          Collection<? extends BoundedWindow> windows,
-          PaneInfo pane) {
-        outputManager.output(mainTag, WindowedValue.of(output, timestamp, windows, pane));
-      }
-
-      @Override
-      public <AdditionalOutputT> void outputWindowedValue(
-          TupleTag<AdditionalOutputT> tag,
-          AdditionalOutputT output,
-          Instant timestamp,
-          Collection<? extends BoundedWindow> windows,
-          PaneInfo pane) {
-        outputManager.output(tag, WindowedValue.of(output, timestamp, windows, pane));
-      }
-    };
-  }
-
   @ProcessElement
   public void processElement(ProcessContext c) throws Exception {
     KeyedWorkItem<K, InputT> keyedWorkItem = c.element();
@@ -130,7 +104,7 @@ public void processElement(ProcessContext c) throws Exception {
                 TriggerStateMachines.stateMachineForTrigger(triggerProto)),
             stateInternals,
             timerInternals,
-            outputWindowedValue(),
+            outputManager.forTag(mainTag),
             sideInputReader,
             reduceFn,
             c.getPipelineOptions());

runners/core-java/src/main/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java

@@ -45,6 +45,7 @@
 import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.util.ValueWithMetadataMultiReceiver;
 import org.apache.beam.sdk.util.WindowedValue;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.PCollectionView;
@@ -72,19 +73,20 @@ public class OutputAndTimeBoundedSplittableProcessElementInvoker<
         InputT, OutputT, RestrictionT, PositionT, WatermarkEstimatorStateT> {
   private final DoFn<InputT, OutputT> fn;
   private final PipelineOptions pipelineOptions;
-  private final OutputWindowedValue<OutputT> output;
+  private final ValueWithMetadataMultiReceiver outputReceiver;
   private final SideInputReader sideInputReader;
   private final ScheduledExecutorService executor;
   private final int maxNumOutputs;
   private final Duration maxDuration;
   private final Supplier<BundleFinalizer> bundleFinalizer;
+  private final TupleTag<OutputT> mainOutputTag;
 
   /**
    * Creates a new invoker from components.
    *
    * @param fn The original {@link DoFn}.
    * @param pipelineOptions {@link PipelineOptions} to include in the {@link DoFn.ProcessContext}.
-   * @param output Hook for outputting from the {@link DoFn.ProcessElement} method.
+   * @param outputReceiver Hook for outputting from the {@link DoFn.ProcessElement} method.
    * @param sideInputReader Hook for accessing side inputs.
    * @param executor Executor on which a checkpoint will be scheduled after the given duration.
    * @param maxNumOutputs Maximum number of outputs, in total over all output tags, after which a
@@ -98,15 +100,17 @@ public class OutputAndTimeBoundedSplittableProcessElementInvoker<
   public OutputAndTimeBoundedSplittableProcessElementInvoker(
       DoFn<InputT, OutputT> fn,
       PipelineOptions pipelineOptions,
-      OutputWindowedValue<OutputT> output,
+      ValueWithMetadataMultiReceiver outputReceiver,
+      TupleTag<OutputT> mainOutputTag,
       SideInputReader sideInputReader,
       ScheduledExecutorService executor,
       int maxNumOutputs,
       Duration maxDuration,
       Supplier<BundleFinalizer> bundleFinalizer) {
     this.fn = fn;
     this.pipelineOptions = pipelineOptions;
-    this.output = output;
+    this.outputReceiver = outputReceiver;
+    this.mainOutputTag = mainOutputTag;
     this.sideInputReader = sideInputReader;
     this.executor = executor;
     this.maxNumOutputs = maxNumOutputs;
@@ -403,7 +407,9 @@ public void outputWindowedValue(
       if (watermarkEstimator instanceof TimestampObservingWatermarkEstimator) {
         ((TimestampObservingWatermarkEstimator) watermarkEstimator).observeTimestamp(timestamp);
       }
-      output.outputWindowedValue(value, timestamp, windows, paneInfo);
+      outputReceiver
+          .forTag(mainOutputTag)
+          .output(WindowedValue.of(value, timestamp, windows, paneInfo));
     }
 
     @Override
@@ -413,7 +419,9 @@ public <T> void output(TupleTag<T> tag, T value) {
 
     @Override
     public <T> void outputWithTimestamp(TupleTag<T> tag, T value, Instant timestamp) {
-      outputWindowedValue(tag, value, timestamp, element.getWindows(), element.getPane());
+      outputReceiver
+          .forTag(tag)
+          .output(WindowedValue.of(value, timestamp, element.getWindows(), element.getPane()));
     }
 
     @Override
@@ -427,7 +435,7 @@ public <T> void outputWindowedValue(
       if (watermarkEstimator instanceof TimestampObservingWatermarkEstimator) {
         ((TimestampObservingWatermarkEstimator) watermarkEstimator).observeTimestamp(timestamp);
       }
-      output.outputWindowedValue(tag, value, timestamp, windows, paneInfo);
+      outputReceiver.forTag(tag).output(WindowedValue.of(value, timestamp, windows, paneInfo));
     }
 
     private void noteOutput() {

runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnRunner.java

@@ -46,6 +46,7 @@
 import org.apache.beam.sdk.transforms.windowing.Window;
 import org.apache.beam.sdk.transforms.windowing.Window.ClosingBehavior;
 import org.apache.beam.sdk.transforms.windowing.WindowFn;
+import org.apache.beam.sdk.util.ValueWithMetadataReceiver;
 import org.apache.beam.sdk.util.WindowTracing;
 import org.apache.beam.sdk.util.WindowedValue;
 import org.apache.beam.sdk.values.KV;
@@ -106,7 +107,7 @@ public class ReduceFnRunner<K, InputT, OutputT, W extends BoundedWindow> {
    */
   private final WindowingStrategy<Object, W> windowingStrategy;
 
-  private final OutputWindowedValue<KV<K, OutputT>> outputter;
+  private final ValueWithMetadataReceiver<KV<K, OutputT>> outputter;
 
   private final StateInternals stateInternals;
 
@@ -214,7 +215,7 @@ public ReduceFnRunner(
       ExecutableTriggerStateMachine triggerStateMachine,
       StateInternals stateInternals,
       TimerInternals timerInternals,
-      OutputWindowedValue<KV<K, OutputT>> outputter,
+      ValueWithMetadataReceiver<KV<K, OutputT>> outputter,
       @Nullable SideInputReader sideInputReader,
       ReduceFn<K, InputT, OutputT, W> reduceFn,
       @Nullable PipelineOptions options) {
@@ -1055,7 +1056,8 @@ private void prefetchOnTrigger(
                 }
 
                 // Output the actual value.
-                outputter.outputWindowedValue(KV.of(key, toOutput), outputTimestamp, windows, pane);
+                outputter.output(
+                    WindowedValue.of(KV.of(key, toOutput), outputTimestamp, windows, pane));
               });
 
       reduceFn.onTrigger(renamedTriggerContext);

runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java

@@ -26,7 +26,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import org.apache.beam.runners.core.DoFnRunners.OutputManager;
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.options.PipelineOptions;
 import org.apache.beam.sdk.schemas.SchemaCoder;
@@ -54,6 +53,7 @@
 import org.apache.beam.sdk.transforms.windowing.PaneInfo;
 import org.apache.beam.sdk.util.SystemDoFnInternal;
 import org.apache.beam.sdk.util.UserCodeException;
+import org.apache.beam.sdk.util.ValueWithMetadataMultiReceiver;
 import org.apache.beam.sdk.util.WindowedValue;
 import org.apache.beam.sdk.values.PCollectionView;
 import org.apache.beam.sdk.values.Row;
@@ -93,7 +93,7 @@ public class SimpleDoFnRunner<InputT, OutputT> implements DoFnRunner<InputT, Out
   private final DoFnInvoker<InputT, OutputT> invoker;
 
   private final SideInputReader sideInputReader;
-  private final OutputManager outputManager;
+  private final ValueWithMetadataMultiReceiver outputManager;
 
   private final TupleTag<OutputT> mainOutputTag;
   /** The set of known output tags. */
@@ -123,7 +123,7 @@ public SimpleDoFnRunner(
       PipelineOptions options,
       DoFn<InputT, OutputT> fn,
       SideInputReader sideInputReader,
-      OutputManager outputManager,
+      ValueWithMetadataMultiReceiver outputManager,
       TupleTag<OutputT> mainOutputTag,
       List<TupleTag<?>> additionalOutputTags,
       StepContext stepContext,
@@ -273,7 +273,7 @@ private void checkTimestamp(Instant elemTimestamp, Instant timestamp) {
 
   private <T> void outputWindowedValue(TupleTag<T> tag, WindowedValue<T> windowedElem) {
     checkArgument(outputTags.contains(tag), "Unknown output tag %s", tag);
-    outputManager.output(tag, windowedElem);
+    outputManager.forTag(tag).output(windowedElem);
   }
 
   /** An {@link DoFnInvoker.ArgumentProvider} for {@link DoFn.StartBundle @StartBundle}. */

runners/core-java/src/test/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvokerTest.java

@@ -27,7 +27,6 @@
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 
-import java.util.Collection;
 import java.util.Collections;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
@@ -38,9 +37,10 @@
 import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
 import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
 import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
-import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
 import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.util.ValueWithMetadataMultiReceiver;
+import org.apache.beam.sdk.util.ValueWithMetadataReceiver;
 import org.apache.beam.sdk.util.WindowedValue;
 import org.apache.beam.sdk.values.TupleTag;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.Uninterruptibles;
@@ -112,22 +112,13 @@ private SplittableProcessElementInvoker<Void, String, OffsetRange, Long, Void>.R
         new OutputAndTimeBoundedSplittableProcessElementInvoker<>(
             fn,
             PipelineOptionsFactory.create(),
-            new OutputWindowedValue<String>() {
+            new ValueWithMetadataMultiReceiver() {
               @Override
-              public void outputWindowedValue(
-                  String output,
-                  Instant timestamp,
-                  Collection<? extends BoundedWindow> windows,
-                  PaneInfo pane) {}
-
-              @Override
-              public <AdditionalOutputT> void outputWindowedValue(
-                  TupleTag<AdditionalOutputT> tag,
-                  AdditionalOutputT output,
-                  Instant timestamp,
-                  Collection<? extends BoundedWindow> windows,
-                  PaneInfo pane) {}
+              public <OutputT> ValueWithMetadataReceiver<OutputT> forTag(TupleTag<OutputT> tag) {
+                return output -> {};
+              }
             },
+            null,
             NullSideInputReader.empty(),
             Executors.newSingleThreadScheduledExecutor(),
             1000,