Extending WindowedValue with ElementMetadata

runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/BatchViewOverrides.java

@@ -63,6 +63,7 @@
 import org.apache.beam.sdk.transforms.windowing.PaneInfo;
 import org.apache.beam.sdk.transforms.windowing.Window;
 import org.apache.beam.sdk.util.CoderUtils;
+import org.apache.beam.sdk.util.ElementMetadata;
 import org.apache.beam.sdk.util.SystemDoFnInternal;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.PCollection;
@@ -1365,16 +1366,28 @@ private static <T> WindowedValue<T> valueInEmptyWindows(T value) {
   private static class ValueInEmptyWindows<T> implements WindowedValue<T> {
 
     private final T value;
+    private final @Nullable ElementMetadata elementMetadata;
 
     private ValueInEmptyWindows(T value) {
       this.value = value;
+      this.elementMetadata = null;
+    }
+
+    private ValueInEmptyWindows(T value, @Nullable ElementMetadata elementMetadata) {
+      this.value = value;
+      this.elementMetadata = elementMetadata;
     }
 
     @Override
     public <NewT> WindowedValue<NewT> withValue(NewT value) {
       return new ValueInEmptyWindows<>(value);
     }
 
+    @Override
+    public WindowedValue<T> withElementMetadata(@Nullable ElementMetadata elementMetadata) {
+      return new ValueInEmptyWindows<>(this.getValue(), elementMetadata);
+    }
+
     @Override
     public T getValue() {
       return value;
@@ -1400,6 +1413,11 @@ public Iterable<WindowedValue<T>> explodeWindows() {
       return Collections.emptyList();
     }
 
+    @Override
+    public @Nullable ElementMetadata getElementMetadata() {
+      return elementMetadata;
+    }
+
     @Override
     public String toString() {
       return MoreObjects.toStringHelper(getClass()).add("value", getValue()).toString();

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java

@@ -109,7 +109,9 @@
 import org.apache.beam.sdk.io.FileSystems;
 import org.apache.beam.sdk.io.gcp.bigquery.BigQuerySinkMetrics;
 import org.apache.beam.sdk.metrics.MetricsEnvironment;
+import org.apache.beam.sdk.transforms.windowing.PaneInfo;
 import org.apache.beam.sdk.util.construction.CoderTranslation;
+import org.apache.beam.sdk.values.WindowedValues;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheStats;
@@ -816,6 +818,9 @@ public static void main(String[] args) throws Exception {
 
     CoderTranslation.verifyModelCodersRegistered();
 
+    WindowedValues.FullWindowedValueCoder.setMetadataSupported();
+    PaneInfo.PaneInfoCoder.setMetadataSupported();
+
     LOG.debug("Creating StreamingDataflowWorker from options: {}", options);
     StreamingDataflowWorker worker = StreamingDataflowWorker.fromOptions(options);
 

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/ValueInEmptyWindows.java

@@ -22,6 +22,7 @@
 import java.util.Objects;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.util.ElementMetadata;
 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;
@@ -39,9 +40,15 @@
  */
 public class ValueInEmptyWindows<T> implements WindowedValue<T> {
   private final T value;
+  private final @Nullable ElementMetadata elementMetadata;
 
   public ValueInEmptyWindows(T value) {
+    this(value, null);
+  }
+
+  public ValueInEmptyWindows(T value, @Nullable ElementMetadata elementMetadata) {
     this.value = value;
+    this.elementMetadata = elementMetadata;
   }
 
   @Override
@@ -54,6 +61,11 @@ public Iterable<WindowedValue<T>> explodeWindows() {
     return Collections.emptyList();
   }
 
+  @Override
+  public @Nullable ElementMetadata getElementMetadata() {
+    return elementMetadata;
+  }
+
   @Override
   public T getValue() {
     return value;
@@ -64,6 +76,11 @@ public <NewT> WindowedValue<NewT> withValue(NewT newValue) {
     return new ValueInEmptyWindows<>(newValue);
   }
 
+  @Override
+  public WindowedValue<T> withElementMetadata(@Nullable ElementMetadata elementMetadata) {
+    return new ValueInEmptyWindows<>(this.getValue(), elementMetadata);
+  }
+
   @Override
   public Instant getTimestamp() {
     return BoundedWindow.TIMESTAMP_MIN_VALUE;

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

@@ -146,10 +146,10 @@ private static byte encodedByte(boolean isFirst, boolean isLast, Timing timing)
     ImmutableMap.Builder<Byte, PaneInfo> decodingBuilder = ImmutableMap.builder();
     for (Timing timing : Timing.values()) {
       long onTimeIndex = timing == Timing.EARLY ? -1 : 0;
-      register(decodingBuilder, new PaneInfo(true, true, timing, 0, onTimeIndex));
-      register(decodingBuilder, new PaneInfo(true, false, timing, 0, onTimeIndex));
-      register(decodingBuilder, new PaneInfo(false, true, timing, -1, onTimeIndex));
-      register(decodingBuilder, new PaneInfo(false, false, timing, -1, onTimeIndex));
+      register(decodingBuilder, new PaneInfo(true, true, timing, 0, onTimeIndex, false));
+      register(decodingBuilder, new PaneInfo(true, false, timing, 0, onTimeIndex, false));
+      register(decodingBuilder, new PaneInfo(false, true, timing, -1, onTimeIndex, false));
+      register(decodingBuilder, new PaneInfo(false, false, timing, -1, onTimeIndex, false));
     }
     BYTE_TO_PANE_INFO = decodingBuilder.build();
   }
@@ -159,7 +159,7 @@ private static void register(ImmutableMap.Builder<Byte, PaneInfo> builder, PaneI
   }
 
   private final byte encodedByte;
-
+  private final boolean containsElementMetadata;
   private final boolean isFirst;
   private final boolean isLast;
   private final Timing timing;
@@ -177,27 +177,46 @@ private static void register(ImmutableMap.Builder<Byte, PaneInfo> builder, PaneI
   public static final PaneInfo ON_TIME_AND_ONLY_FIRING =
       PaneInfo.createPane(true, true, Timing.ON_TIME, 0, 0);
 
-  private PaneInfo(boolean isFirst, boolean isLast, Timing timing, long index, long onTimeIndex) {
+  private PaneInfo(
+      boolean isFirst,
+      boolean isLast,
+      Timing timing,
+      long index,
+      long onTimeIndex,
+      boolean containsElementMetadata) {
     this.encodedByte = encodedByte(isFirst, isLast, timing);
     this.isFirst = isFirst;
     this.isLast = isLast;
     this.timing = timing;
     this.index = index;
     this.nonSpeculativeIndex = onTimeIndex;
+    this.containsElementMetadata = containsElementMetadata;
   }
 
   public static PaneInfo createPane(boolean isFirst, boolean isLast, Timing timing) {
     checkArgument(isFirst, "Indices must be provided for non-first pane info.");
     return createPane(isFirst, isLast, timing, 0, timing == Timing.EARLY ? -1 : 0);
   }
 
+  /** Factory method to create a {@link PaneInfo} with the specified parameters. */
   /** Factory method to create a {@link PaneInfo} with the specified parameters. */
   public static PaneInfo createPane(
       boolean isFirst, boolean isLast, Timing timing, long index, long onTimeIndex) {
+    return createPane(isFirst, isLast, timing, index, onTimeIndex, false);
+  }
+
+  /** Factory method to create a {@link PaneInfo} with the specified parameters. */
+  public static PaneInfo createPane(
+      boolean isFirst,
+      boolean isLast,
+      Timing timing,
+      long index,
+      long onTimeIndex,
+      boolean containsElementMetadata) {
     if (isFirst || timing == Timing.UNKNOWN) {
       return checkNotNull(BYTE_TO_PANE_INFO.get(encodedByte(isFirst, isLast, timing)));
     } else {
-      return new PaneInfo(isFirst, isLast, timing, index, onTimeIndex);
+      return new PaneInfo(isFirst, isLast, timing, index, onTimeIndex, containsElementMetadata);
     }
   }
 
@@ -219,6 +238,15 @@ public boolean isFirst() {
     return isFirst;
   }
 
+  public boolean isElementMetadata() {
+    return containsElementMetadata;
+  }
+
+  public PaneInfo withElementMetadata(boolean elementMetadata) {
+    return new PaneInfo(
+        this.isFirst, this.isLast, this.timing, index, nonSpeculativeIndex, elementMetadata);
+  }
+
   /** Return true if this is the last pane that will be produced in the associated window. */
   public boolean isLast() {
     return isLast;
@@ -295,6 +323,13 @@ public String toString() {
 
   /** A Coder for encoding PaneInfo instances. */
   public static class PaneInfoCoder extends AtomicCoder<PaneInfo> {
+    private static final byte ELEMENT_METADATA_MASK = (byte) 0x80;
+    protected static boolean metadataSupported = false;
+
+    public static void setMetadataSupported() {
+      metadataSupported = true;
+    }
+
     private enum Encoding {
       FIRST,
       ONE_INDEX,
@@ -337,16 +372,17 @@ private PaneInfoCoder() {}
     public void encode(PaneInfo value, final OutputStream outStream)
         throws CoderException, IOException {
       Encoding encoding = chooseEncoding(value);
+      byte elementMetadata = value.containsElementMetadata ? ELEMENT_METADATA_MASK : 0x00;
       switch (chooseEncoding(value)) {
         case FIRST:
-          outStream.write(value.encodedByte);
+          outStream.write(value.encodedByte | elementMetadata);
           break;
         case ONE_INDEX:
-          outStream.write(value.encodedByte | encoding.tag);
+          outStream.write(value.encodedByte | encoding.tag | elementMetadata);
           VarInt.encode(value.index, outStream);
           break;
         case TWO_INDICES:
-          outStream.write(value.encodedByte | encoding.tag);
+          outStream.write(value.encodedByte | encoding.tag | elementMetadata);
           VarInt.encode(value.index, outStream);
           VarInt.encode(value.nonSpeculativeIndex, outStream);
           break;
@@ -360,9 +396,10 @@ public PaneInfo decode(final InputStream inStream) throws CoderException, IOExce
       byte keyAndTag = (byte) inStream.read();
       PaneInfo base = Preconditions.checkNotNull(BYTE_TO_PANE_INFO.get((byte) (keyAndTag & 0x0F)));
       long index, onTimeIndex;
-      switch (Encoding.fromTag(keyAndTag)) {
+      boolean elementMetadata = (keyAndTag & ELEMENT_METADATA_MASK) != 0;
+      switch (Encoding.fromTag((byte) (keyAndTag & ~ELEMENT_METADATA_MASK))) {
         case FIRST:
-          return base;
+          return base.withElementMetadata(elementMetadata);
         case ONE_INDEX:
           index = VarInt.decodeLong(inStream);
           onTimeIndex = base.timing == Timing.EARLY ? -1 : index;
@@ -374,7 +411,8 @@ public PaneInfo decode(final InputStream inStream) throws CoderException, IOExce
         default:
           throw new CoderException("Unknown encoding " + (keyAndTag & 0xF0));
       }
-      return new PaneInfo(base.isFirst, base.isLast, base.timing, index, onTimeIndex);
+      return new PaneInfo(
+          base.isFirst, base.isLast, base.timing, index, onTimeIndex, elementMetadata);
     }
 
     @Override

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

@@ -0,0 +1,28 @@
+/*
+ * 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.sdk.util;
+
+import com.google.auto.value.AutoValue;
+
+@AutoValue
+public abstract class ElementMetadata {
+
+  public static ElementMetadata create() {
+    return new AutoValue_ElementMetadata();
+  }
+}

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

@@ -495,6 +495,7 @@ public static Set<String> getJavaCapabilities() {
     capabilities.add(BeamUrns.getUrn(StandardProtocols.Enum.DATA_SAMPLING));
     capabilities.add(BeamUrns.getUrn(StandardProtocols.Enum.SDK_CONSUMING_RECEIVED_DATA));
     capabilities.add(BeamUrns.getUrn(StandardProtocols.Enum.ORDERED_LIST_STATE));
+    capabilities.add(BeamUrns.getUrn(StandardProtocols.Enum.ELEMENT_METADATA));
     return capabilities.build();
   }
 

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

@@ -20,6 +20,8 @@
 import java.util.Collection;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.util.ElementMetadata;
+import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.Instant;
 
 /**
@@ -34,6 +36,9 @@ public interface WindowedValue<T> {
   /** The timestamp of this value in event time. */
   Instant getTimestamp();
 
+  @Nullable
+  ElementMetadata getElementMetadata();
+
   /** Returns the windows of this {@code WindowedValue}. */
   Collection<? extends BoundedWindow> getWindows();
 
@@ -57,4 +62,6 @@ default PaneInfo getPaneInfo() {
    * value.
    */
   <OtherT> WindowedValue<OtherT> withValue(OtherT value);
+
+  WindowedValue<T> withElementMetadata(@Nullable ElementMetadata elementMetadata);
 }