This is an automated email from the ASF dual-hosted git repository.
shunping pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 03db1a07096 Fix DataflowOutputCounter calculation for
ValueInEmptyWindows (#39487)
03db1a07096 is described below
commit 03db1a070961eb8fd5a04d9124a3b4423919672b
Author: Shunping Huang <[email protected]>
AuthorDate: Thu Jul 30 17:47:46 2026 -0400
Fix DataflowOutputCounter calculation for ValueInEmptyWindows (#39487)
* Fix DataflowOutputCounter calculation for ValueInEmptyWindows
When processing shuffle or streaming data in Dataflow Legacy Runner
(e.g., from GroupingShuffleReader or WindowingWindmillReader),
KeyedWorkItems are wrapped inside a ValueInEmptyWindows (windows.size() ==
0).
Previously, DataflowOutputCounter.update() counted these as 1 element.
This caused inaccurate element counts because:
1. A KeyedWorkItem can contain multiple elements.
2. Elements may belong to multiple windows and need to be fanned out
accordingly.
3. KeyedWorkItems containing only timers were incorrectly incrementing
element counters.
* Address non keyedworkitems
* Spotless
* Add elementWindowsIterable to only decode window metadata and use it in
DataflowOutputCounter
* Use the elementWindowsIterable in ReduceFnRunner
* Refactor: separate batch and streaming DataflowOutputCounter
implementations
* Minor change on tests.
* Address reviewer comments
* Spotless
* Remove unnecessary comments
* Address comments
---
.../core/GroupAlsoByWindowViaWindowSetNewDoFn.java | 2 +-
.../apache/beam/runners/core/KeyedWorkItem.java | 9 ++
.../apache/beam/runners/core/ReduceFnRunner.java | 18 +++-
.../dataflow/worker/DataflowOutputCounter.java | 68 ++++++++++---
.../worker/IntrinsicMapTaskExecutorFactory.java | 11 ++-
.../dataflow/worker/SimpleParDoFnHelpers.java | 5 +-
.../StreamingGroupAlsoByWindowViaWindowSetFn.java | 2 +-
.../dataflow/worker/WindmillKeyedWorkItem.java | 26 ++++-
.../dataflow/worker/DataflowOutputCounterTest.java | 108 +++++++++++++++++++++
.../IntrinsicMapTaskExecutorFactoryTest.java | 14 ++-
10 files changed, 234 insertions(+), 29 deletions(-)
diff --git
a/runners/core-java/src/main/java/org/apache/beam/runners/core/GroupAlsoByWindowViaWindowSetNewDoFn.java
b/runners/core-java/src/main/java/org/apache/beam/runners/core/GroupAlsoByWindowViaWindowSetNewDoFn.java
index f242c7d1000..349e109930f 100644
---
a/runners/core-java/src/main/java/org/apache/beam/runners/core/GroupAlsoByWindowViaWindowSetNewDoFn.java
+++
b/runners/core-java/src/main/java/org/apache/beam/runners/core/GroupAlsoByWindowViaWindowSetNewDoFn.java
@@ -109,7 +109,7 @@ public class GroupAlsoByWindowViaWindowSetNewDoFn<
reduceFn,
c.getPipelineOptions());
- reduceFnRunner.processElements(keyedWorkItem.elementsIterable());
+ reduceFnRunner.processElements(keyedWorkItem);
reduceFnRunner.onTimers(keyedWorkItem.timersIterable());
reduceFnRunner.persist();
}
diff --git
a/runners/core-java/src/main/java/org/apache/beam/runners/core/KeyedWorkItem.java
b/runners/core-java/src/main/java/org/apache/beam/runners/core/KeyedWorkItem.java
index 4901c5cbed5..2be8b079030 100644
---
a/runners/core-java/src/main/java/org/apache/beam/runners/core/KeyedWorkItem.java
+++
b/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();
+ }
}
diff --git
a/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnRunner.java
b/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnRunner.java
index 7fe3b711aa0..e49c858393f 100644
---
a/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnRunner.java
+++
b/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnRunner.java
@@ -60,6 +60,7 @@ import
org.apache.beam.sdk.values.WindowingStrategy.AccumulationMode;
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 @@ public class ReduceFnRunner<K, InputT, OutputT, W extends
BoundedWindow> {
* 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 class ReduceFnRunner<K, InputT, OutputT, W extends
BoundedWindow> {
}
/** 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()) {
diff --git
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowOutputCounter.java
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowOutputCounter.java
index 7c5859a9d32..1a927c03c61 100644
---
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowOutputCounter.java
+++
b/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 @@ import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.Vi
@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 class DataflowOutputCounter implements
ElementCounter {
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);
diff --git
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/IntrinsicMapTaskExecutorFactory.java
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/IntrinsicMapTaskExecutorFactory.java
index d3f2aacc74d..3ea29787eb3 100644
---
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/IntrinsicMapTaskExecutorFactory.java
+++
b/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.Coder;
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 class IntrinsicMapTaskExecutorFactory implements
DataflowMapTaskExecutorF
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 @@ public class IntrinsicMapTaskExecutorFactory implements
DataflowMapTaskExecutorF
* 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,7 +357,7 @@ public class IntrinsicMapTaskExecutorFactory implements
DataflowMapTaskExecutorF
CloudObjects.coderFromCloudObject(CloudObject.fromSpec(cloudOutput.getCodec()));
ElementCounter outputCounter =
- new DataflowOutputCounter(
+ DataflowOutputCounter.create(
cloudOutput.getName(),
new ElementByteSizeObservableCoder<>(coder),
counterFactory,
@@ -363,7 +365,8 @@ public class IntrinsicMapTaskExecutorFactory implements
DataflowMapTaskExecutorF
stageName,
cloudOutput.getOriginalName(),
cloudOutput.getSystemName(),
- cloudOutput.getName()));
+ cloudOutput.getName()),
+ isStreaming);
outputReceiver.addOutputCounter(outputCounter);
return OutputReceiverNode.create(outputReceiver, coder,
input.getPcollectionId());
diff --git
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFnHelpers.java
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFnHelpers.java
index 964cf2323d5..15bfba9bbc4 100644
---
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFnHelpers.java
+++
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFnHelpers.java
@@ -190,9 +190,10 @@ class SimpleParDoFnHelpers<InputT, OutputT, W extends
BoundedWindow> {
// 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;
diff --git
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowViaWindowSetFn.java
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowViaWindowSetFn.java
index ec36644d1e6..a183df19b6e 100644
---
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowViaWindowSetFn.java
+++
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowViaWindowSetFn.java
@@ -93,7 +93,7 @@ public class StreamingGroupAlsoByWindowViaWindowSetFn<K,
InputT, OutputT, W exte
reduceFn,
options);
- reduceFnRunner.processElements(keyedWorkItem.elementsIterable());
+ reduceFnRunner.processElements(keyedWorkItem);
reduceFnRunner.onTimers(keyedWorkItem.timersIterable());
reduceFnRunner.persist();
}
diff --git
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItem.java
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItem.java
index 82116e0b2d8..ff5be071ede 100644
---
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItem.java
+++
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItem.java
@@ -140,6 +140,16 @@ public class WindmillKeyedWorkItem<K, ElemT> implements
KeyedWorkItem<K, ElemT>
}
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 =
@@ -162,8 +172,11 @@ public class WindmillKeyedWorkItem<K, ElemT> implements
KeyedWorkItem<K, ElemT>
valueKind =
WindmillValueKindHelper.fromProto(elementMetadata.getValueKind());
openTelemetryContext =
WindmillOpenTelemetryContextPropagator.read(elementMetadata);
}
- 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,
@@ -187,6 +200,15 @@ public class WindmillKeyedWorkItem<K, ElemT> implements
KeyedWorkItem<K, ElemT>
}
}
+ @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() {
diff --git
a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/DataflowOutputCounterTest.java
b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/DataflowOutputCounterTest.java
new file mode 100644
index 00000000000..b5c49ee639b
--- /dev/null
+++
b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/DataflowOutputCounterTest.java
@@ -0,0 +1,108 @@
+/*
+ * 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.runners.dataflow.worker;
+
+import static org.junit.Assert.assertEquals;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
+
+import java.util.Arrays;
+import org.apache.beam.runners.core.KeyedWorkItem;
+import org.apache.beam.runners.dataflow.worker.counters.CounterName;
+import org.apache.beam.runners.dataflow.worker.counters.CounterSet;
+import org.apache.beam.runners.dataflow.worker.counters.NameContext;
+import org.apache.beam.runners.dataflow.worker.util.ValueInEmptyWindows;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.WindowedValue;
+import org.apache.beam.sdk.values.WindowedValues;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests for {@link DataflowOutputCounter}. */
+@RunWith(JUnit4.class)
+public class DataflowOutputCounterTest {
+ private static final String OUTPUT_NAME = "test_output";
+ private CounterSet counterSet;
+ private NameContext nameContext;
+
+ @Before
+ public void setUp() {
+ counterSet = new CounterSet();
+ nameContext = NameContext.create("stage", "original", "system",
OUTPUT_NAME);
+ }
+
+ @Test
+ public void testBatchOutputCounterWithEmptyWindows() throws Exception {
+ DataflowOutputCounter batchCounter =
+ DataflowOutputCounter.create(OUTPUT_NAME, counterSet, nameContext,
false);
+
+ ValueInEmptyWindows<KV<String, String>> shuffleValue =
+ new ValueInEmptyWindows<>(KV.of("key", "value"));
+ batchCounter.update(shuffleValue);
+
+ long elementCount =
+ (Long)
+ counterSet
+ .getExistingCounter(
+
CounterName.named(DataflowOutputCounter.getElementCounterName(OUTPUT_NAME)))
+ .getAggregate();
+ assertEquals(1L, elementCount);
+ }
+
+ @Test
+ public void testStreamingOutputCounterWithKeyedWorkItem() throws Exception {
+ DataflowOutputCounter streamingCounter =
+ DataflowOutputCounter.create(OUTPUT_NAME, counterSet, nameContext,
true);
+
+ KeyedWorkItem<String, String> kwi = mock(KeyedWorkItem.class);
+ WindowedValue<String> element1 = WindowedValues.valueInGlobalWindow("v1");
+ WindowedValue<String> element2 = WindowedValues.valueInGlobalWindow("v2");
+ doReturn(Arrays.asList(element1,
element2)).when(kwi).elementWindowsIterable();
+
+ ValueInEmptyWindows<KeyedWorkItem<String, String>> streamingValue =
+ new ValueInEmptyWindows<>(kwi);
+ streamingCounter.update(streamingValue);
+
+ long elementCount =
+ (Long)
+ counterSet
+ .getExistingCounter(
+
CounterName.named(DataflowOutputCounter.getElementCounterName(OUTPUT_NAME)))
+ .getAggregate();
+ assertEquals(2L, elementCount);
+ }
+
+ @Test
+ public void testStandardWindowedValueCounting() throws Exception {
+ DataflowOutputCounter counter =
+ DataflowOutputCounter.create(OUTPUT_NAME, counterSet, nameContext,
false);
+
+ WindowedValue<String> standardValue =
WindowedValues.valueInGlobalWindow("v1");
+ counter.update(standardValue);
+
+ long elementCount =
+ (Long)
+ counterSet
+ .getExistingCounter(
+
CounterName.named(DataflowOutputCounter.getElementCounterName(OUTPUT_NAME)))
+ .getAggregate();
+ assertEquals(1L, elementCount);
+ }
+}
diff --git
a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/IntrinsicMapTaskExecutorFactoryTest.java
b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/IntrinsicMapTaskExecutorFactoryTest.java
index 3443ae0022b..d3a424758f6 100644
---
a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/IntrinsicMapTaskExecutorFactoryTest.java
+++
b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/IntrinsicMapTaskExecutorFactoryTest.java
@@ -330,7 +330,8 @@ public class IntrinsicMapTaskExecutorFactoryTest {
when(network.successors(instructionNode))
.thenReturn(
ImmutableSet.<Node>of(
-
IntrinsicMapTaskExecutorFactory.createOutputReceiversTransform(STAGE,
counterSet)
+ IntrinsicMapTaskExecutorFactory.createOutputReceiversTransform(
+ STAGE, counterSet, false)
.apply(
InstructionOutputNode.create(
instructionNode.getParallelInstruction().getOutputs().get(0),
@@ -535,7 +536,7 @@ public class IntrinsicMapTaskExecutorFactoryTest {
ExecutionLocation.UNKNOWN);
Node outputReceiverNode =
- IntrinsicMapTaskExecutorFactory.createOutputReceiversTransform(STAGE,
counterSet)
+ IntrinsicMapTaskExecutorFactory.createOutputReceiversTransform(STAGE,
counterSet, false)
.apply(
InstructionOutputNode.create(
instructionNode.getParallelInstruction().getOutputs().get(0), PCOLLECTION_ID));
@@ -614,7 +615,8 @@ public class IntrinsicMapTaskExecutorFactoryTest {
when(network.successors(instructionNode))
.thenReturn(
ImmutableSet.<Node>of(
-
IntrinsicMapTaskExecutorFactory.createOutputReceiversTransform(STAGE,
counterSet)
+ IntrinsicMapTaskExecutorFactory.createOutputReceiversTransform(
+ STAGE, counterSet, false)
.apply(
InstructionOutputNode.create(
instructionNode.getParallelInstruction().getOutputs().get(0),
@@ -669,7 +671,8 @@ public class IntrinsicMapTaskExecutorFactoryTest {
when(network.successors(instructionNode))
.thenReturn(
ImmutableSet.<Node>of(
-
IntrinsicMapTaskExecutorFactory.createOutputReceiversTransform(STAGE,
counterSet)
+ IntrinsicMapTaskExecutorFactory.createOutputReceiversTransform(
+ STAGE, counterSet, false)
.apply(
InstructionOutputNode.create(
instructionNode.getParallelInstruction().getOutputs().get(0),
@@ -750,7 +753,8 @@ public class IntrinsicMapTaskExecutorFactoryTest {
when(network.successors(instructionNode))
.thenReturn(
ImmutableSet.<Node>of(
-
IntrinsicMapTaskExecutorFactory.createOutputReceiversTransform(STAGE,
counterSet)
+ IntrinsicMapTaskExecutorFactory.createOutputReceiversTransform(
+ STAGE, counterSet, false)
.apply(
InstructionOutputNode.create(
instructionNode.getParallelInstruction().getOutputs().get(0),
