Fix DataflowOutputCounter calculation for ValueInEmptyWindows

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

@@ -109,7 +109,7 @@ public void processElement(ProcessContext c) throws Exception {
             reduceFn,
             c.getPipelineOptions());
 
-    reduceFnRunner.processElements(keyedWorkItem.elementsIterable());
+    reduceFnRunner.processElements(keyedWorkItem);
     reduceFnRunner.onTimers(keyedWorkItem.timersIterable());
     reduceFnRunner.persist();
   }

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

@@ -35,4 +35,13 @@ public interface KeyedWorkItem<K, ElemT> {
 
   /** Returns an iterable containing the elements. */
   Iterable<WindowedValue<ElemT>> elementsIterable();
+
+  /**
+   * Returns an iterable containing windowed values without guaranteeing element payload decoding.
+   * Useful for lightweight inspection of windowing metadata without payload deserialization
+   * overhead.
+   */
+  default Iterable<WindowedValue<?>> elementWindowsIterable() {
+    return (Iterable) elementsIterable();
+  }
 }

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

@@ -60,6 +60,7 @@
 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.collect.FluentIterable;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.Duration;
 import org.joda.time.Instant;
@@ -361,13 +362,24 @@ private Collection<W> windowsThatShouldFire(Set<W> windows) throws Exception {
    *       setting holds, and invoking {@link ReduceFn#onTrigger}.
    * </ol>
    */
+  public void processElements(KeyedWorkItem<?, InputT> keyedWorkItem) throws Exception {
+    processElementsInternal(
+        keyedWorkItem.elementWindowsIterable(), keyedWorkItem.elementsIterable());
+  }
+
   public void processElements(Iterable<WindowedValue<InputT>> values) throws Exception {
-    if (!values.iterator().hasNext()) {
+    processElementsInternal(values, values);
+  }
+
+  private void processElementsInternal(
+      Iterable<? extends WindowedValue<?>> elementWindows, Iterable<WindowedValue<InputT>> values)
+      throws Exception {
+    if (Iterables.isEmpty(elementWindows)) {
       return;
     }
 
     // Determine all the windows for elements.
-    Set<W> windows = collectWindows(values);
+    Set<W> windows = collectWindows(elementWindows);
     // If an incoming element introduces a new window, attempt to merge it into an existing
     // window eagerly.
     Map<W, W> windowToMergeResult = mergeWindows(windows);
@@ -426,7 +438,7 @@ public void persist() {
   }
 
   /** Extract the windows associated with the values. */
-  private Set<W> collectWindows(Iterable<WindowedValue<InputT>> values) throws Exception {
+  private Set<W> collectWindows(Iterable<? extends WindowedValue<?>> values) throws Exception {
     Set<W> windows = new HashSet<>();
     for (WindowedValue<?> value : values) {
       for (BoundedWindow untypedWindow : value.getWindows()) {

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

@@ -18,12 +18,14 @@
 package org.apache.beam.runners.dataflow.worker;
 
 import org.apache.beam.runners.core.ElementByteSizeObservable;
+import org.apache.beam.runners.core.KeyedWorkItem;
 import org.apache.beam.runners.dataflow.worker.counters.Counter;
 import org.apache.beam.runners.dataflow.worker.counters.CounterFactory;
 import org.apache.beam.runners.dataflow.worker.counters.CounterName;
 import org.apache.beam.runners.dataflow.worker.counters.NameContext;
 import org.apache.beam.runners.dataflow.worker.util.common.worker.ElementCounter;
 import org.apache.beam.runners.dataflow.worker.util.common.worker.OutputObjectAndByteCounter;
+import org.apache.beam.sdk.annotations.Internal;
 import org.apache.beam.sdk.values.WindowedValue;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
 
@@ -33,6 +35,7 @@
 @SuppressWarnings({
   "nullness" // TODO(https://github.com/apache/beam/issues/20497)
 })
+@Internal
 public class DataflowOutputCounter implements ElementCounter {
   /** Number of logical element and single window pairs that were processed. */
   private static final String ELEMENT_COUNTER_NAME = "-ElementCount";
@@ -41,20 +44,36 @@ public class DataflowOutputCounter implements ElementCounter {
 
   private OutputObjectAndByteCounter objectAndByteCounter;
   private Counter<Long, ?> elementCount;
+  private final boolean isStreaming;
 
-  public DataflowOutputCounter(
-      String outputName, CounterFactory counterFactory, NameContext nameContext) {
-    this(outputName, null, counterFactory, nameContext);
+  public static DataflowOutputCounter create(
+      String outputName,
+      ElementByteSizeObservable<?> elementByteSizeObservable,
+      CounterFactory counterFactory,
+      NameContext nameContext,
+      boolean isStreaming) {
+    return new DataflowOutputCounter(
+        outputName, elementByteSizeObservable, counterFactory, nameContext, isStreaming);
+  }
+
+  public static DataflowOutputCounter create(
+      String outputName,
+      CounterFactory counterFactory,
+      NameContext nameContext,
+      boolean isStreaming) {
+    return new DataflowOutputCounter(outputName, null, counterFactory, nameContext, isStreaming);
   }
 
-  public DataflowOutputCounter(
+  private DataflowOutputCounter(
       String outputName,
       ElementByteSizeObservable<?> elementByteSizeObservable,
       CounterFactory counterFactory,
-      NameContext nameContext) {
-    objectAndByteCounter =
+      NameContext nameContext,
+      boolean isStreaming) {
+    this.isStreaming = isStreaming;
+    this.objectAndByteCounter =
         new OutputObjectAndByteCounter(elementByteSizeObservable, counterFactory, nameContext);
-    objectAndByteCounter.countMeanByte(outputName + MEAN_BYTE_COUNTER_NAME);
+    this.objectAndByteCounter.countMeanByte(outputName + MEAN_BYTE_COUNTER_NAME);
     createElementCounter(counterFactory, outputName + ELEMENT_COUNTER_NAME);
   }
 
@@ -63,15 +82,42 @@ public void update(Object elem) throws Exception {
     objectAndByteCounter.update(elem);
     long windowsSize = ((WindowedValue<?>) elem).getWindows().size();
     if (windowsSize == 0) {
-      // GroupingShuffleReader produces ValueInEmptyWindows.
-      // For now, we count the element at least once to keep the current counter
-      // behavior.
-      elementCount.addValue(1L);
+      updateEmptyWindows((WindowedValue<?>) elem);
     } else {
+      // Standard WindowedValue.
       elementCount.addValue(windowsSize);
     }
   }
 
+  private void updateEmptyWindows(WindowedValue<?> elem) {
+    if (isStreaming) {
+      Object value = elem.getValue();
+      if (value instanceof KeyedWorkItem<?, ?>) {
+        // KeyedWorkItem wrapped in ValueInEmptyWindows
+        // (e.g. WindowingWindmillReader for Streaming GBK)
+        KeyedWorkItem<?, ?> keyedWorkItem = (KeyedWorkItem<?, ?>) value;
+        long totalElementCount = 0;
+        // Iterate through elementWindowsIterable and ignore timers in KeyedWorkItem.
+        // Uses lightweight metadata-only iteration without payload deserialization overhead.
+        for (WindowedValue<?> element : keyedWorkItem.elementWindowsIterable()) {
+          long elementWindowsSize = element.getWindows().size();
+          // Fan out for windows.
+          totalElementCount += (elementWindowsSize == 0 ? 1L : elementWindowsSize);
+        }
+        elementCount.addValue(totalElementCount);
+      } else {
+        // NOTE: in streaming mode, this should not normally happen.
+        // Counting as 1 element serves as a fallback to maintain counter behavior without failing
+        // execution.
+        elementCount.addValue(1L);
+      }
+    } else {
+      // Non-KeyedWorkItem wrapped in ValueInEmptyWindows
+      // (e.g. GroupingShuffleReader KV output for Batch GBK)
+      elementCount.addValue(1L);
+    }
+  }
+
   @Override
   public void finishLazyUpdate(Object elem) {
     objectAndByteCounter.finishLazyUpdate(elem);

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

@@ -63,6 +63,7 @@
 import org.apache.beam.sdk.coders.KvCoder;
 import org.apache.beam.sdk.fn.IdGenerator;
 import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.options.StreamingOptions;
 import org.apache.beam.sdk.util.common.ElementByteSizeObserver;
 import org.apache.beam.sdk.values.TupleTag;
 import org.apache.beam.sdk.values.WindowedValues.WindowedValueCoder;
@@ -102,8 +103,9 @@ public DataflowMapTaskExecutor create(
       IdGenerator idGenerator) {
 
     // Swap out all the InstructionOutput nodes with OutputReceiver nodes
+    boolean isStreaming = options.as(StreamingOptions.class).isStreaming();
     Networks.replaceDirectedNetworkNodes(
-        network, createOutputReceiversTransform(stageName, counterSet));
+        network, createOutputReceiversTransform(stageName, counterSet, isStreaming));
 
     // Swap out all the ParallelInstruction nodes with Operation nodes. While updating the network,
     // we keep track of
@@ -345,7 +347,7 @@ OperationNode createFlattenOperation(
    * Returns a function which can convert {@link InstructionOutput}s into {@link OutputReceiver}s.
    */
   static Function<Node, Node> createOutputReceiversTransform(
-      final String stageName, final CounterFactory counterFactory) {
+      final String stageName, final CounterFactory counterFactory, final boolean isStreaming) {
     return new TypeSafeNodeFunction<InstructionOutputNode>(InstructionOutputNode.class) {
       @Override
       public Node typedApply(InstructionOutputNode input) {
@@ -355,15 +357,16 @@ public Node typedApply(InstructionOutputNode input) {
             CloudObjects.coderFromCloudObject(CloudObject.fromSpec(cloudOutput.getCodec()));
 
         ElementCounter outputCounter =
-            new DataflowOutputCounter(
+            DataflowOutputCounter.create(
                 cloudOutput.getName(),
                 new ElementByteSizeObservableCoder<>(coder),
                 counterFactory,
                 NameContext.create(
                     stageName,
                     cloudOutput.getOriginalName(),
                     cloudOutput.getSystemName(),
-                    cloudOutput.getName()));
+                    cloudOutput.getName()),
+                isStreaming);
         outputReceiver.addOutputCounter(outputCounter);
 
         return OutputReceiverNode.create(outputReceiver, coder, input.getPcollectionId());

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

@@ -190,9 +190,10 @@ public <TagT> void output(TupleTag<TagT> tag, WindowedValue<TagT> output) {
               // doesn't today.)
               OutputReceiver undeclaredReceiver = new OutputReceiver();
 
+              boolean isStreaming = options.as(StreamingOptions.class).isStreaming();
               ElementCounter outputCounter =
-                  new DataflowOutputCounter(
-                      outputName, counterFactory, stepContext.getNameContext());
+                  DataflowOutputCounter.create(
+                      outputName, counterFactory, stepContext.getNameContext(), isStreaming);
               undeclaredReceiver.addOutputCounter(outputCounter);
               undeclaredOutputs.put(tag, undeclaredReceiver);
               receiver = undeclaredReceiver;

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

@@ -93,7 +93,7 @@ public void processElement(
             reduceFn,
             options);
 
-    reduceFnRunner.processElements(keyedWorkItem.elementsIterable());
+    reduceFnRunner.processElements(keyedWorkItem);
     reduceFnRunner.onTimers(keyedWorkItem.timersIterable());
     reduceFnRunner.persist();
   }

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

@@ -139,6 +139,16 @@ public Iterable<TimerData> timersIterable() {
   }
 
   private @Nullable WindowedValue<ElemT> parseElem(Windmill.Message message) {
+    return parseElemInternal(message, true);
+  }
+
+  private @Nullable WindowedValue<?> parseElemWindowOnly(Windmill.Message message) {
+    return parseElemInternal(message, false);
+  }
+
+  @SuppressWarnings("nullness")
+  private @Nullable WindowedValue<ElemT> parseElemInternal(
+      Windmill.Message message, boolean parseValue) {
     try {
       Instant timestamp = WindmillTimeUtils.windmillToHarnessTimestamp(message.getTimestamp());
       Collection<? extends BoundedWindow> windows =
@@ -159,8 +169,11 @@ public Iterable<TimerData> timersIterable() {
                 : CausedByDrain.NORMAL;
         valueKind = WindmillValueKindHelper.fromProto(elementMetadata.getValueKind());
       }
-      InputStream inputStream = message.getData().newInput();
-      ElemT value = valueCoder.decode(inputStream, Coder.Context.OUTER);
+      ElemT value = null;
+      if (parseValue) {
+        InputStream inputStream = message.getData().newInput();
+        value = valueCoder.decode(inputStream, Coder.Context.OUTER);
+      }
       return WindowedValues.of(
           value,
           timestamp,
@@ -184,6 +197,15 @@ public Iterable<TimerData> timersIterable() {
     }
   }
 
+  @Override
+  @SuppressWarnings("nullness")
+  public Iterable<WindowedValue<?>> elementWindowsIterable() {
+    return FluentIterable.from(workItem.getMessageBundlesList())
+        .transformAndConcat(Windmill.InputMessageBundle::getMessagesList)
+        .transform(this::parseElemWindowOnly)
+        .filter(Objects::nonNull);
+  }
+
   @Override
   @SuppressWarnings("nullness")
   public Iterable<WindowedValue<ElemT>> elementsIterable() {