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 f242c7d10003..349e109930f0 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 void processElement(ProcessContext c) throws Exception { 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 4901c5cbed5b..2be8b0790301 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 { /** Returns an iterable containing the elements. */ Iterable> 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> 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 7fe3b711aa0a..e49c858393f4 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.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 windowsThatShouldFire(Set windows) throws Exception { * setting holds, and invoking {@link ReduceFn#onTrigger}. * */ + public void processElements(KeyedWorkItem keyedWorkItem) throws Exception { + processElementsInternal( + keyedWorkItem.elementWindowsIterable(), keyedWorkItem.elementsIterable()); + } + public void processElements(Iterable> values) throws Exception { - if (!values.iterator().hasNext()) { + processElementsInternal(values, values); + } + + private void processElementsInternal( + Iterable> elementWindows, Iterable> values) + throws Exception { + if (Iterables.isEmpty(elementWindows)) { return; } // Determine all the windows for elements. - Set windows = collectWindows(values); + Set windows = collectWindows(elementWindows); // If an incoming element introduces a new window, attempt to merge it into an existing // window eagerly. Map windowToMergeResult = mergeWindows(windows); @@ -426,7 +438,7 @@ public void persist() { } /** Extract the windows associated with the values. */ - private Set collectWindows(Iterable> values) throws Exception { + private Set collectWindows(Iterable> values) throws Exception { Set 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 7c5859a9d324..1a927c03c61c 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 @@ @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 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); 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 d3f2aacc74d0..3ea29787eb3b 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.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 createOutputReceiversTransform( - final String stageName, final CounterFactory counterFactory) { + final String stageName, final CounterFactory counterFactory, final boolean isStreaming) { return new TypeSafeNodeFunction(InstructionOutputNode.class) { @Override public Node typedApply(InstructionOutputNode input) { @@ -355,7 +357,7 @@ public Node typedApply(InstructionOutputNode input) { CloudObjects.coderFromCloudObject(CloudObject.fromSpec(cloudOutput.getCodec())); ElementCounter outputCounter = - new DataflowOutputCounter( + DataflowOutputCounter.create( cloudOutput.getName(), new ElementByteSizeObservableCoder<>(coder), counterFactory, @@ -363,7 +365,8 @@ public Node typedApply(InstructionOutputNode input) { 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 964cf2323d51..15bfba9bbc49 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 @@ public void output(TupleTag tag, WindowedValue 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; 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 ec36644d1e68..a183df19b6e7 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 void processElement( 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 c4c0b6ed92d3..7ab18867fdd8 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 @@ -139,6 +139,16 @@ public Iterable timersIterable() { } private @Nullable WindowedValue parseElem(Windmill.Message message) { + return parseElemInternal(message, true); + } + + private @Nullable WindowedValue parseElemWindowOnly(Windmill.Message message) { + return parseElemInternal(message, false); + } + + @SuppressWarnings("nullness") + private @Nullable WindowedValue parseElemInternal( + Windmill.Message message, boolean parseValue) { try { Instant timestamp = WindmillTimeUtils.windmillToHarnessTimestamp(message.getTimestamp()); Collection windows = @@ -159,8 +169,11 @@ public Iterable 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 timersIterable() { } } + @Override + @SuppressWarnings("nullness") + public Iterable> elementWindowsIterable() { + return FluentIterable.from(workItem.getMessageBundlesList()) + .transformAndConcat(Windmill.InputMessageBundle::getMessagesList) + .transform(this::parseElemWindowOnly) + .filter(Objects::nonNull); + } + @Override @SuppressWarnings("nullness") public Iterable> 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 000000000000..b5c49ee639b5 --- /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> 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 kwi = mock(KeyedWorkItem.class); + WindowedValue element1 = WindowedValues.valueInGlobalWindow("v1"); + WindowedValue element2 = WindowedValues.valueInGlobalWindow("v2"); + doReturn(Arrays.asList(element1, element2)).when(kwi).elementWindowsIterable(); + + ValueInEmptyWindows> 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 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 3443ae0022bc..d3a424758f66 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 void testCreateReadOperation() throws Exception { when(network.successors(instructionNode)) .thenReturn( ImmutableSet.of( - IntrinsicMapTaskExecutorFactory.createOutputReceiversTransform(STAGE, counterSet) + IntrinsicMapTaskExecutorFactory.createOutputReceiversTransform( + STAGE, counterSet, false) .apply( InstructionOutputNode.create( instructionNode.getParallelInstruction().getOutputs().get(0), @@ -535,7 +536,7 @@ public void testCreateParDoOperation() throws Exception { 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 void testCreatePartialGroupByKeyOperation() throws Exception { when(network.successors(instructionNode)) .thenReturn( ImmutableSet.of( - IntrinsicMapTaskExecutorFactory.createOutputReceiversTransform(STAGE, counterSet) + IntrinsicMapTaskExecutorFactory.createOutputReceiversTransform( + STAGE, counterSet, false) .apply( InstructionOutputNode.create( instructionNode.getParallelInstruction().getOutputs().get(0), @@ -669,7 +671,8 @@ public void testCreatePartialGroupByKeyOperationWithCombine() throws Exception { when(network.successors(instructionNode)) .thenReturn( ImmutableSet.of( - IntrinsicMapTaskExecutorFactory.createOutputReceiversTransform(STAGE, counterSet) + IntrinsicMapTaskExecutorFactory.createOutputReceiversTransform( + STAGE, counterSet, false) .apply( InstructionOutputNode.create( instructionNode.getParallelInstruction().getOutputs().get(0), @@ -750,7 +753,8 @@ public void testCreateFlattenOperation() throws Exception { when(network.successors(instructionNode)) .thenReturn( ImmutableSet.of( - IntrinsicMapTaskExecutorFactory.createOutputReceiversTransform(STAGE, counterSet) + IntrinsicMapTaskExecutorFactory.createOutputReceiversTransform( + STAGE, counterSet, false) .apply( InstructionOutputNode.create( instructionNode.getParallelInstruction().getOutputs().get(0),