WIP: add OutputBuilder

Author: kennknowles

buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy

@@ -1182,6 +1182,7 @@ class BeamModulePlugin implements Plugin<Project> {
 
       List<String> skipDefRegexes = []
       skipDefRegexes << "AutoValue_.*"
+      skipDefRegexes << "AutoBuilder_.*"
       skipDefRegexes << "AutoOneOf_.*"
       skipDefRegexes << ".*\\.jmh_generated\\..*"
       skipDefRegexes += configuration.generatedClassPatterns
@@ -1275,7 +1276,8 @@ class BeamModulePlugin implements Plugin<Project> {
         '**/org/apache/beam/gradle/**',
         '**/org/apache/beam/model/**',
         '**/org/apache/beam/runners/dataflow/worker/windmill/**',
-        '**/AutoValue_*'
+        '**/AutoValue_*',
+        '**/AutoBuilder_*',
       ]
 
       def jacocoEnabled = project.hasProperty('enableJacocoReport')

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

@@ -45,6 +45,7 @@
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.transforms.windowing.PaneInfo;
 import org.apache.beam.sdk.transforms.windowing.Window;
+import org.apache.beam.sdk.values.OutputBuilder;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.sdk.values.PCollectionView;
 import org.apache.beam.sdk.values.Row;
@@ -391,17 +392,27 @@ public TypeDescriptor<OutputT> getOutputTypeDescriptor() {
 
   /** Receives values of the given type. */
   public interface OutputReceiver<T> {
-    void output(T output);
+    OutputBuilder<T> builder();
 
-    void outputWithTimestamp(T output, Instant timestamp);
+    default void output(T value) {
+      builder().setValue(value).output();
+    }
+
+    default void outputWithTimestamp(T value, Instant timestamp) {
+      builder().setValue(value).setTimestamp(timestamp).output();
+    }
 
     default void outputWindowedValue(
-        T output,
+        T value,
         Instant timestamp,
         Collection<? extends BoundedWindow> windows,
         PaneInfo paneInfo) {
-      throw new UnsupportedOperationException(
-          String.format("Not implemented: %s.outputWindowedValue", this.getClass().getName()));
+      builder()
+          .setValue(value)
+          .setTimestamp(timestamp)
+          .setWindows(windows)
+          .setPaneInfo(paneInfo)
+          .output();
     }
   }
 

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

@@ -17,6 +17,7 @@
  */
 package org.apache.beam.sdk.transforms;
 
+import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull;
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
 
@@ -29,6 +30,8 @@
 import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.util.WindowedValue;
+import org.apache.beam.sdk.values.OutputBuilder;
 import org.apache.beam.sdk.values.Row;
 import org.apache.beam.sdk.values.TupleTag;
 import org.checkerframework.checker.nullness.qual.Nullable;
@@ -40,40 +43,127 @@
   "nullness" // TODO(https://github.com/apache/beam/issues/20497)
 })
 public class DoFnOutputReceivers {
-  private static class RowOutputReceiver<T> implements OutputReceiver<Row> {
-    WindowedContextOutputReceiver<T> outputReceiver;
-    SchemaCoder<T> schemaCoder;
 
-    public RowOutputReceiver(
-        DoFn<?, ?>.WindowedContext context,
-        @Nullable TupleTag<T> outputTag,
-        SchemaCoder<T> schemaCoder) {
-      outputReceiver = new WindowedContextOutputReceiver<>(context, outputTag);
-      this.schemaCoder = checkNotNull(schemaCoder);
+  /**
+   * OutputBuilder implementations may extend this class to simplify wrapping of another
+   * OutputBuilder.
+   *
+   * <p>Because a common cause of delegation is to adjust the type of output received, extenders
+   * <i>must</i> implement {#link setValue} and {#link getValue}.
+   */
+  public abstract static class DelegatingOutputBuilder<InnerT, OuterT>
+      implements OutputBuilder<OuterT> {
+
+    private final OutputBuilder<InnerT> delegate;
+
+    public DelegatingOutputBuilder(OutputBuilder<InnerT> delegate) {
+      this.delegate = delegate;
     }
 
+    protected OutputBuilder<InnerT> getDelegate() {
+      return delegate;
+    }
+
+    @Override
+    public abstract OutputBuilder<OuterT> setValue(OuterT value);
+
     @Override
-    public void output(Row output) {
-      outputReceiver.output(schemaCoder.getFromRowFunction().apply(output));
+    public OutputBuilder<OuterT> setTimestamp(Instant timestamp) {
+      delegate.setTimestamp(timestamp);
+      return this;
     }
 
     @Override
-    public void outputWithTimestamp(Row output, Instant timestamp) {
-      outputReceiver.outputWithTimestamp(schemaCoder.getFromRowFunction().apply(output), timestamp);
+    public OutputBuilder<OuterT> setWindow(BoundedWindow window) {
+      delegate.setWindow(window);
+      return this;
     }
 
     @Override
-    public void outputWindowedValue(
-        Row output,
-        Instant timestamp,
-        Collection<? extends BoundedWindow> windows,
-        PaneInfo paneInfo) {
-      outputReceiver.outputWindowedValue(
-          schemaCoder.getFromRowFunction().apply(output), timestamp, windows, paneInfo);
+    public OutputBuilder<OuterT> setWindows(Collection<? extends BoundedWindow> windows) {
+      delegate.setWindows(windows);
+      return this;
+    }
+
+    @Override
+    public OutputBuilder<OuterT> setPaneInfo(PaneInfo paneInfo) {
+      delegate.setPaneInfo(paneInfo);
+      return this;
+    }
+
+    @Override
+    public abstract OuterT getValue();
+
+    @Override
+    public Instant getTimestamp() {
+      return delegate.getTimestamp();
+    }
+
+    @Override
+    public Collection<? extends BoundedWindow> getWindows() {
+      return delegate.getWindows();
+    }
+
+    @Override
+    public PaneInfo getPaneInfo() {
+      return delegate.getPaneInfo();
+    }
+
+    @Override
+    public void output() {
+      delegate.output();
     }
   }
 
-  private static class WindowedContextOutputReceiver<T> implements OutputReceiver<T> {
+  private static class RowOutputBuilder<T> extends DelegatingOutputBuilder<T, Row>
+      implements OutputBuilder<Row> {
+
+    private @Nullable Row currentRow;
+    private final SchemaCoder<T> schemaCoder;
+
+    RowOutputBuilder(OutputBuilder<T> delegate, SchemaCoder<T> schemaCoder) {
+      super(delegate);
+      this.schemaCoder = schemaCoder;
+    }
+
+    @Override
+    public OutputBuilder<Row> setValue(Row value) {
+      currentRow = value;
+      getDelegate().setValue(schemaCoder.getFromRowFunction().apply(value));
+      return (OutputBuilder<Row>) this;
+    }
+
+    @Override
+    public Row getValue() {
+      checkStateNotNull(currentRow, "getValue() called before value set");
+      return currentRow;
+    }
+  }
+
+  private static class RowOutputReceiver<T> implements OutputReceiver<Row> {
+    WindowedContextOutputReceiver<T> outputReceiver;
+    SchemaCoder<T> schemaCoder;
+
+    @Override
+    public OutputBuilder<Row> builder() {
+      return new RowOutputBuilder<>(outputReceiver.builder(), schemaCoder);
+    }
+
+    public RowOutputReceiver(
+        DoFn<?, ?>.WindowedContext context,
+        @Nullable TupleTag<T> outputTag,
+        SchemaCoder<T> schemaCoder) {
+      outputReceiver = new WindowedContextOutputReceiver<>(context, outputTag);
+      this.schemaCoder = checkNotNull(schemaCoder);
+    }
+  }
+
+  /**
+   * OutputReceiver that delegates all its core functionality to DoFn.WindowedContext which predates
+   * OutputReceiver and has most of the same methods.
+   */
+  private static class WindowedContextOutputReceiver<T>
+      implements OutputReceiver<T>, WindowedValue.Outputter<T> {
     DoFn<?, ?>.WindowedContext context;
     @Nullable TupleTag<T> outputTag;
 
@@ -84,34 +174,26 @@ public WindowedContextOutputReceiver(
     }
 
     @Override
-    public void output(T output) {
-      if (outputTag != null) {
-        context.output(outputTag, output);
-      } else {
-        ((DoFn<?, T>.WindowedContext) context).output(output);
-      }
-    }
-
-    @Override
-    public void outputWithTimestamp(T output, Instant timestamp) {
-      if (outputTag != null) {
-        context.outputWithTimestamp(outputTag, output, timestamp);
-      } else {
-        ((DoFn<?, T>.WindowedContext) context).outputWithTimestamp(output, timestamp);
-      }
+    public OutputBuilder<T> builder() {
+      return WindowedValue.builder(this);
     }
 
     @Override
-    public void outputWindowedValue(
-        T output,
-        Instant timestamp,
-        Collection<? extends BoundedWindow> windows,
-        PaneInfo paneInfo) {
+    public void output(WindowedValue.Builder<T> outputBuilder) {
       if (outputTag != null) {
-        context.outputWindowedValue(outputTag, output, timestamp, windows, paneInfo);
+        context.outputWindowedValue(
+            outputTag,
+            outputBuilder.getValue(),
+            outputBuilder.getTimestamp(),
+            outputBuilder.getWindows(),
+            outputBuilder.getPaneInfo());
       } else {
         ((DoFn<?, T>.WindowedContext) context)
-            .outputWindowedValue(output, timestamp, windows, paneInfo);
+            .outputWindowedValue(
+                outputBuilder.getValue(),
+                outputBuilder.getTimestamp(),
+                outputBuilder.getWindows(),
+                outputBuilder.getPaneInfo());
       }
     }
   }

sdks/java/core/src/main/java/org/apache/beam/sdk/util/WindowedValue.java

@@ -17,9 +17,11 @@
  */
 package org.apache.beam.sdk.util;
 
+import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull;
 
+import com.google.auto.value.AutoBuilder;
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
@@ -46,8 +48,11 @@
 import org.apache.beam.sdk.transforms.windowing.PaneInfo;
 import org.apache.beam.sdk.transforms.windowing.PaneInfo.PaneInfoCoder;
 import org.apache.beam.sdk.util.common.ElementByteSizeObserver;
+import org.apache.beam.sdk.values.OutputBuilder;
+import org.apache.beam.vendor.grpc.v1p69p0.com.google.common.collect.Iterables;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.Instant;
 
@@ -59,16 +64,92 @@
 @Internal
 public abstract class WindowedValue<T> {
 
+  public interface Outputter<T> {
+    void output(Builder<T> outputBuilder);
+  }
+
+  /** Create a typical {@link OutputBuilder} that */
+  public static <T> OutputBuilder<T> builder(Outputter<T> outputter) {
+    return new AutoBuilder_WindowedValue_Builder<T>().setOutputter(outputter);
+  }
+
+  // HACK: this should not be public but for temporary ease of demonstration in FnApiDoFnRunner
+  @AutoBuilder(callMethod = "of", ofClass = WindowedValue.class)
+  public abstract static class Builder<T> implements OutputBuilder<T> {
+
+    private @MonotonicNonNull Outputter<T> outputter;
+
+    @Override
+    public abstract OutputBuilder<T> setValue(T value);
+
+    @Override
+    public abstract OutputBuilder<T> setTimestamp(Instant timestamp);
+
+    @Override
+    public abstract OutputBuilder<T> setWindows(Collection<? extends BoundedWindow> windows);
+
+    @Override
+    public abstract OutputBuilder<T> setPaneInfo(PaneInfo pane);
+
+    @Override
+    public OutputBuilder<T> setWindow(BoundedWindow window) {
+      return setWindows(Collections.singleton(window));
+    }
+
+    public Builder<T> setOutputter(Outputter<T> outputter) {
+      this.outputter = outputter;
+      return this;
+    }
+
+    abstract WindowedValue<T> autoBuild();
+
+    @Override
+    public void output() {
+      checkStateNotNull(outputter, "Cannot output until Outputter set").output(this);
+    }
+
+    @Override
+    public abstract T getValue();
+
+    @Override
+    public abstract Instant getTimestamp();
+
+    @Override
+    public abstract Collection<? extends BoundedWindow> getWindows();
+
+    @Override
+    public abstract PaneInfo getPaneInfo();
+
+    public WindowedValue<T> build() {
+      if (getWindows().size() > 1) {
+        return autoBuild();
+      }
+
+      BoundedWindow window = Iterables.getOnlyElement(getWindows());
+
+      if (!GlobalWindow.INSTANCE.equals(window)) {
+        return new TimestampedValueInSingleWindow<>(
+            getValue(), getTimestamp(), window, getPaneInfo());
+      }
+
+      if (BoundedWindow.TIMESTAMP_MIN_VALUE.equals(getTimestamp())) {
+        return new ValueInGlobalWindow<>(getValue(), getPaneInfo());
+      }
+
+      return new TimestampedValueInGlobalWindow<>(getValue(), getTimestamp(), getPaneInfo());
+    };
+  }
+
   /** Returns a {@code WindowedValue} with the given value, timestamp, and windows. */
   public static <T> WindowedValue<T> of(
-      T value, Instant timestamp, Collection<? extends BoundedWindow> windows, PaneInfo pane) {
-    checkArgument(pane != null, "WindowedValue requires PaneInfo, but it was null");
+      T value, Instant timestamp, Collection<? extends BoundedWindow> windows, PaneInfo paneInfo) {
+    checkArgument(paneInfo != null, "WindowedValue requires PaneInfo, but it was null");
     checkArgument(windows.size() > 0, "WindowedValue requires windows, but there were none");
 
     if (windows.size() == 1) {
-      return of(value, timestamp, windows.iterator().next(), pane);
+      return of(value, timestamp, windows.iterator().next(), paneInfo);
     } else {
-      return new TimestampedValueInMultipleWindows<>(value, timestamp, windows, pane);
+      return new TimestampedValueInMultipleWindows<>(value, timestamp, windows, paneInfo);
     }
   }