From 827847d05ced1e956a298aa5cc97fddd5947a5a5 Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Fri, 16 Oct 2026 12:41:41 -0700 Subject: [PATCH 1/9] Adds bytes read reporting for SDF SplittableProcessFn ParDos. --- ...oundedSplittableProcessElementInvoker.java | 2 +- ...edSplittableProcessElementInvokerTest.java | 70 ++++ .../IntrinsicMapTaskExecutorFactory.java | 23 +- .../worker/SplittableProcessFnFactory.java | 4 +- .../worker/StreamingModeExecutionContext.java | 38 ++- .../util/common/worker/MapTaskExecutor.java | 15 + .../util/common/worker/ParDoOperation.java | 21 ++ .../ComputationWorkExecutorFactory.java | 41 +++ .../IntrinsicMapTaskExecutorFactoryTest.java | 92 +++++ .../StreamingModeExecutionContextTest.java | 149 ++++++++ .../common/worker/MapTaskExecutorTest.java | 34 ++ .../ComputationWorkExecutorFactoryTest.java | 323 ++++++++++++++++++ 12 files changed, 792 insertions(+), 20 deletions(-) create mode 100644 runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactoryTest.java diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java index 59d33c596c8a..df1136a5d821 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java @@ -182,7 +182,7 @@ public TimeDomain timeDomain(DoFn doFn) { @Override public OutputReceiver outputReceiver(DoFn doFn) { return DoFnOutputReceivers.windowedReceiver( - processContext, OutputBuilderSuppliers.supplierForElement(element), null); + processContext, OutputBuilderSuppliers.supplierForElement(element), mainOutputTag); } @Override diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvokerTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvokerTest.java index 6ff40fae91dd..32f21a32cb18 100644 --- a/runners/core-java/src/test/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvokerTest.java +++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvokerTest.java @@ -27,7 +27,9 @@ import static org.junit.Assert.assertNull; import static org.junit.Assert.assertTrue; +import java.util.ArrayList; import java.util.Collections; +import java.util.List; import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; import org.apache.beam.sdk.io.range.OffsetRange; @@ -40,6 +42,7 @@ import org.apache.beam.sdk.transforms.windowing.GlobalWindow; import org.apache.beam.sdk.transforms.windowing.PaneInfo; import org.apache.beam.sdk.util.WindowedValueMultiReceiver; +import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.sdk.values.WindowedValues; @@ -287,4 +290,71 @@ public void testBacklogBytesWhenDone() throws Exception { assertEquals(0.0, res.getBacklogBytes(), 0.001); assertNull(res.getResidualRestriction()); } + + private static class OutputReceiverFn extends DoFn { + @ProcessElement + public ProcessContinuation process( + OutputReceiver receiver, RestrictionTracker tracker) { + for (long i = tracker.currentRestriction().getFrom(); tracker.tryClaim(i); ++i) { + receiver.output(String.valueOf(i)); + } + return stop(); + } + + @GetInitialRestriction + public OffsetRange getInitialRestriction(@SuppressWarnings("unused") @Element Void element) { + return new OffsetRange(0, 3); + } + } + + @Test + public void testInvokeProcessElementWithOutputReceiverUsesMainOutputTag() throws Exception { + OutputReceiverFn fn = new OutputReceiverFn(); + TupleTag mainOutputTag = new TupleTag<>("mainOutputTag"); + List, String>> outputs = new ArrayList<>(); + + SplittableProcessElementInvoker invoker = + new OutputAndTimeBoundedSplittableProcessElementInvoker<>( + fn, + PipelineOptionsFactory.create(), + new WindowedValueMultiReceiver() { + @Override + public void output(TupleTag tag, WindowedValue output) { + outputs.add(KV.of(tag, (String) output.getValue())); + } + }, + mainOutputTag, + NullSideInputReader.empty(), + Executors.newSingleThreadScheduledExecutor(), + 1000, + Duration.standardMinutes(3), + () -> { + throw new UnsupportedOperationException("BundleFinalizer not configured for test."); + }); + + SplittableProcessElementInvoker.Result res = + invoker.invokeProcessElement( + DoFnInvokers.invokerFor(fn), + WindowedValues.of(null, Instant.now(), GlobalWindow.INSTANCE, PaneInfo.NO_FIRING), + new OffsetRangeTracker(new OffsetRange(0, 3)), + new WatermarkEstimator() { + @Override + public Instant currentWatermark() { + return GlobalWindow.TIMESTAMP_MIN_VALUE; + } + + @Override + public Void getState() { + return null; + } + }, + Collections.emptyMap()); + + assertFalse(res.getContinuation().shouldResume()); + assertNull(res.getResidualRestriction()); + assertEquals(3, outputs.size()); + assertEquals(KV.of(mainOutputTag, "0"), outputs.get(0)); + assertEquals(KV.of(mainOutputTag, "1"), outputs.get(1)); + assertEquals(KV.of(mainOutputTag, "2"), outputs.get(2)); + } } 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 3ea29787eb3b..7589a5036b58 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 @@ -141,11 +141,9 @@ public DataflowMapTaskExecutor create( Iterables.filter(Networks.topologicalOrder(network), OperationNode.class)) { topoSortedOperations.add(node.getOperation()); - for (Node predecessor : + for (OutputReceiverNode predecessor : Iterables.filter(network.predecessors(node), OutputReceiverNode.class)) { - ((OutputReceiverNode) predecessor) - .getOutputReceiver() - .addOutput((Receiver) node.getOperation()); + predecessor.getOutputReceiver().addOutput((Receiver) node.getOperation()); } } @@ -248,7 +246,7 @@ OperationNode createWriteOperation( return OperationNode.create(WriteOperation.create(sink, EMPTY_OUTPUT_RECEIVER_ARRAY, context)); } - private ParDoFnFactory parDoFnFactory = new DefaultParDoFnFactory(); + private final ParDoFnFactory parDoFnFactory = new DefaultParDoFnFactory(); private OperationNode createParDoOperation( Network network, @@ -272,18 +270,29 @@ private OperationNode createParDoOperation( } successorOffset += 1; } + CloudObject cloudUserFn = CloudObject.fromSpec(parDo.getUserFn()); ParDoFn fn = parDoFnFactory.create( options, - CloudObject.fromSpec(parDo.getUserFn()), + cloudUserFn, parDo.getSideInputs(), mainOutputTag, outputTagsToReceiverIndicesBuilder.build(), executionContext, operationContext); + // An SDF's SplittableProcessFn behaves like a source, so we need to report its bytes read even + // though it's not a read operator for autoscaling to work properly. + Coder bytesReadCoder = null; + if ("SplittableProcessFn".equals(cloudUserFn.getClassName())) { + Node firstOutput = Iterables.getFirst(network.successors(node), null); + if (firstOutput instanceof OutputReceiverNode) { + bytesReadCoder = ((OutputReceiverNode) firstOutput).getCoder(); + } + } OutputReceiver[] receivers = getOutputReceivers(network, node); - return OperationNode.create(new ParDoOperation(fn, receivers, operationContext)); + return OperationNode.create( + new ParDoOperation(fn, receivers, operationContext, bytesReadCoder)); } private static TupleTag tupleTag(MultiOutputInfo multiOutputInfo) { diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SplittableProcessFnFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SplittableProcessFnFactory.java index b55d73cf5793..42aaae3d95ba 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SplittableProcessFnFactory.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SplittableProcessFnFactory.java @@ -151,11 +151,11 @@ public DoFnRunner>, OutputT> crea mainOutputTag, sideInputReader, ses.get(), - // Commit at least once every 10 seconds or 10k records. This keeps the watermark + // Commit at least once every 5 seconds or 10k records. This keeps the watermark // advancing smoothly, and ensures that not too much work will have to be reprocessed // in the event of a crash. 10000, - Duration.standardSeconds(10), + Duration.standardSeconds(5), stepContext::bundleFinalizer)); processFn.setBacklogBytesCallback(userStepContext::setBacklogBytes); DoFnRunner>, OutputT> simpleRunner = diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java index c7036aed0bae..ada5dd1a4af1 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java @@ -68,6 +68,8 @@ import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor; 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.runners.dataflow.worker.util.common.worker.ParDoOperation; +import org.apache.beam.runners.dataflow.worker.util.common.worker.ReadOperation; import org.apache.beam.runners.dataflow.worker.util.common.worker.WorkExecutor; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.GlobalDataId; @@ -796,19 +798,35 @@ private Windmill.WorkItemCommitRequest.Builder buildWorkItemTruncationRequestBui return outputBuilder; } - private final long computeSourceBytesProcessed(String sourceBytesCounterName) { + @VisibleForTesting + long computeSourceBytesProcessed(String sourceBytesCounterName) { if (!(workExecutor instanceof DataflowMapTaskExecutor)) { return 0L; } - HashMap counters = - ((DataflowMapTaskExecutor) workExecutor) - .getReadOperation() - .receivers[0] - .getOutputCounters(); - - return Optional.ofNullable(counters.get(sourceBytesCounterName)) - .map(counter -> ((OutputObjectAndByteCounter) counter).getByteCount().getAndReset()) - .orElse(0L); + DataflowMapTaskExecutor mapTaskExecutor = (DataflowMapTaskExecutor) workExecutor; + long totalBytes = 0L; + ReadOperation readOperation = mapTaskExecutor.getReadOperation(); + if (readOperation != null + && readOperation.receivers != null + && readOperation.receivers.length > 0) { + HashMap readOpCounters = + readOperation.receivers[0].getOutputCounters(); + totalBytes += + Optional.ofNullable(readOpCounters.get(sourceBytesCounterName)) + .map(counter -> ((OutputObjectAndByteCounter) counter).getByteCount().getAndReset()) + .orElse(0L); + } + for (ParDoOperation parDoOp : mapTaskExecutor.getReadReportingParDoOperations()) { + if (parDoOp.receivers != null && parDoOp.receivers.length > 0) { + HashMap splittableProcessOpCounters = + parDoOp.receivers[0].getOutputCounters(); + totalBytes += + Optional.ofNullable(splittableProcessOpCounters.get(sourceBytesCounterName)) + .map(counter -> ((OutputObjectAndByteCounter) counter).getByteCount().getAndReset()) + .orElse(0L); + } + } + return totalBytes; } public boolean advance() throws CoderException { diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutor.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutor.java index 90fc76276940..68c94dc89f3c 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutor.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutor.java @@ -145,6 +145,21 @@ public ReadOperation getReadOperation() { return (ReadOperation) readOperation; } + public List getReadReportingParDoOperations() { + List reportingParDoOperations = Lists.newArrayList(); + if (operations != null) { + for (Operation op : operations) { + if (op instanceof ParDoOperation) { + ParDoOperation parDoOp = (ParDoOperation) op; + if (parDoOp.shouldReportBytesRead()) { + reportingParDoOperations.add(parDoOp); + } + } + } + } + return reportingParDoOperations; + } + public boolean supportsRestart() { for (Operation op : operations) { if (!op.supportsRestart()) { diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoOperation.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoOperation.java index a814b90cae21..224a3b39577a 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoOperation.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoOperation.java @@ -18,16 +18,29 @@ package org.apache.beam.runners.dataflow.worker.util.common.worker; import java.io.Closeable; +import org.apache.beam.sdk.coders.Coder; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; import org.checkerframework.checker.nullness.qual.Nullable; /** A ParDo mapping function. */ public class ParDoOperation extends ReceivingOperation { private final ParDoFn fn; + // If present, this ParDo should report the amount of bytes read using this coder to serialize + // elements. + private final @Nullable Coder bytesReadCoder; public ParDoOperation(ParDoFn fn, OutputReceiver[] outputReceivers, OperationContext context) { + this(fn, outputReceivers, context, null); + } + + public ParDoOperation( + ParDoFn fn, + OutputReceiver[] outputReceivers, + OperationContext context, + @Nullable Coder bytesReadCoder) { super(outputReceivers, context); this.fn = fn; + this.bytesReadCoder = bytesReadCoder; } @Override @@ -82,4 +95,12 @@ public boolean supportsRestart() { public ParDoFn getFn() throws Exception { return fn; } + + public boolean shouldReportBytesRead() { + return bytesReadCoder != null; + } + + public @Nullable Coder getBytesReadCoder() { + return bytesReadCoder; + } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java index b51512252e37..622371d9fcaa 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java @@ -54,6 +54,7 @@ import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputStateFetcherFactory; import org.apache.beam.runners.dataflow.worker.util.common.worker.MapTaskExecutor; import org.apache.beam.runners.dataflow.worker.util.common.worker.OutputObjectAndByteCounter; +import org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoOperation; import org.apache.beam.runners.dataflow.worker.util.common.worker.ReadOperation; import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillStateCache; import org.apache.beam.runners.dataflow.worker.windmill.work.processing.failures.FailureTracker; @@ -185,6 +186,30 @@ private static void trackAutoscalingBytesRead( .countBytes(counterName)); } + private static void trackAutoscalingParDoBytesRead( + MapTask mapTask, + Nodes.ParallelInstructionNode readNode, + Coder splittableProcessCoder, + ParDoOperation splittableParDoOp, + MapTaskExecutor mapTaskExecutor, + String counterName) { + NameContext nameContext = + NameContext.create( + mapTask.getStageName(), + readNode.getParallelInstruction().getOriginalName(), + readNode.getParallelInstruction().getSystemName(), + readNode.getParallelInstruction().getName()); + splittableParDoOp.receivers[0].addOutputCounter( + counterName, + new OutputObjectAndByteCounter( + new IntrinsicMapTaskExecutorFactory.ElementByteSizeObservableCoder<>( + splittableProcessCoder), + mapTaskExecutor.getOutputCounters(), + nameContext) + .setSamplingPeriod(100) + .countBytes(counterName)); + } + private static ReadOperation getValidatedReadOperation(MapTaskExecutor mapTaskExecutor) { ReadOperation readOperation = mapTaskExecutor.getReadOperation(); // Disable progress updates since its results are unused for streaming @@ -234,6 +259,22 @@ ComputationWorkExecutor createComputationWorkExecutor( readOperation, mapTaskExecutor, computationState.sourceBytesProcessCounterName()); + } else { + for (ParDoOperation parDoOp : mapTaskExecutor.getReadReportingParDoOperations()) { + Coder bytesReadCoder = parDoOp.getBytesReadCoder(); + if (parDoOp.receivers != null + && parDoOp.receivers.length > 0 + && parDoOp.receivers[0] != null + && bytesReadCoder != null) { + trackAutoscalingParDoBytesRead( + mapTask, + readNode, + bytesReadCoder, + parDoOp, + mapTaskExecutor, + computationState.sourceBytesProcessCounterName()); + } + } } ComputationWorkExecutor.Builder executionStateBuilder = 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 85973372e419..13c1a46704f0 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 @@ -17,6 +17,7 @@ */ package org.apache.beam.runners.dataflow.worker; +import static org.apache.beam.runners.dataflow.util.Structs.addObject; import static org.apache.beam.runners.dataflow.util.Structs.addString; import static org.apache.beam.runners.dataflow.worker.DataflowOutputCounter.getElementCounterName; import static org.apache.beam.runners.dataflow.worker.DataflowOutputCounter.getMeanByteCounterName; @@ -75,6 +76,7 @@ import org.apache.beam.runners.dataflow.worker.graph.Nodes.Node; import org.apache.beam.runners.dataflow.worker.graph.Nodes.OperationNode; import org.apache.beam.runners.dataflow.worker.graph.Nodes.ParallelInstructionNode; +import org.apache.beam.runners.dataflow.worker.util.WorkerPropertyNames; import org.apache.beam.runners.dataflow.worker.util.common.worker.FlattenOperation; import org.apache.beam.runners.dataflow.worker.util.common.worker.Operation; import org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoOperation; @@ -87,14 +89,17 @@ import org.apache.beam.sdk.coders.IterableCoder; import org.apache.beam.sdk.coders.KvCoder; import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.coders.VoidCoder; import org.apache.beam.sdk.extensions.gcp.util.Transport; import org.apache.beam.sdk.fn.IdGenerator; import org.apache.beam.sdk.fn.IdGenerators; +import org.apache.beam.sdk.io.range.OffsetRange; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.DoFnSchemaInformation; import org.apache.beam.sdk.transforms.Sum; +import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker; import org.apache.beam.sdk.transforms.windowing.IntervalWindow.IntervalWindowCoder; import org.apache.beam.sdk.util.AppliedCombineFn; import org.apache.beam.sdk.util.DoFnInfo; @@ -567,6 +572,93 @@ public void testCreateParDoOperation() throws Exception { assertEquals(1, parDoOperation.receivers.length); assertEquals(0, parDoOperation.receivers[0].getReceiverCount()); assertEquals(Operation.InitializationState.UNSTARTED, parDoOperation.initializationState); + assertFalse(parDoOperation.shouldReportBytesRead()); + assertNull(parDoOperation.getBytesReadCoder()); + } + + static class TestSplittableDoFn extends DoFn { + @ProcessElement + public void processElement(ProcessContext c, RestrictionTracker tracker) { + if (tracker.tryClaim(tracker.currentRestriction().getFrom())) { + c.output(c.element()); + } + } + + @GetInitialRestriction + public OffsetRange getInitialRestriction(@SuppressWarnings("unused") @Element String element) { + return new OffsetRange(0, 1); + } + } + + @Test + public void testCreateSplittableProcessFnParDoOperation() throws Exception { + int producerIndex = 1; + int producerOutputNum = 2; + + BatchModeExecutionContext context = + BatchModeExecutionContext.forTesting(options, counterSet, "testStage"); + TupleTag mainOutput = new TupleTag<>(PropertyNames.OUTPUT); + String serializedFn = + StringUtils.byteArrayToJsonString( + SerializableUtils.serializeToByteArray( + DoFnInfo.forFn( + new TestSplittableDoFn(), + WindowingStrategy.globalDefault(), + null /* side input views */, + StringUtf8Coder.of() /* input coder */, + Collections.singletonMap(mainOutput, StringUtf8Coder.of()), + mainOutput /* main output id */, + DoFnSchemaInformation.create(), + Collections.emptyMap()))); + + CloudObject cloudUserFn = CloudObject.forClassName("SplittableProcessFn"); + addString(cloudUserFn, PropertyNames.SERIALIZED_FN, serializedFn); + addObject( + cloudUserFn, + WorkerPropertyNames.RESTRICTION_CODER, + CloudObjects.asCloudObject( + KvCoder.of(OffsetRange.Coder.of(), VoidCoder.of()), /* sdkComponents= */ null)); + + ParallelInstruction instruction = + createParDoInstruction( + producerIndex, producerOutputNum, "SplittableParDo", "", new TestSplittableDoFn()); + instruction.getParDo().setUserFn(cloudUserFn); + + ParallelInstructionNode instructionNode = + ParallelInstructionNode.create(instruction, ExecutionLocation.UNKNOWN); + + Node outputReceiverNode = + IntrinsicMapTaskExecutorFactory.createOutputReceiversTransform(STAGE, counterSet, false) + .apply( + InstructionOutputNode.create( + instructionNode.getParallelInstruction().getOutputs().get(0), PCOLLECTION_ID)); + + when(network.successors(instructionNode)).thenReturn(ImmutableSet.of(outputReceiverNode)); + when(network.outDegree(instructionNode)).thenReturn(1); + when(network.edgesConnecting(instructionNode, outputReceiverNode)) + .thenReturn( + ImmutableSet.of( + MultiOutputInfoEdge.create( + instructionNode + .getParallelInstruction() + .getParDo() + .getMultiOutputInfos() + .get(0)))); + + ArrayList createdOperations = new ArrayList<>(); + Node operationNode = + mapTaskExecutorFactory + .createOperationTransformForParallelInstructionNodes( + STAGE, network, options, readerRegistry, sinkRegistry, context, createdOperations) + .apply(instructionNode); + assertThat(operationNode, instanceOf(OperationNode.class)); + assertThat(((OperationNode) operationNode).getOperation(), instanceOf(ParDoOperation.class)); + ParDoOperation parDoOperation = (ParDoOperation) ((OperationNode) operationNode).getOperation(); + assertThat(createdOperations, contains(parDoOperation)); + + assertTrue(parDoOperation.shouldReportBytesRead()); + assertEquals( + WindowedValues.getValueOnlyCoder(StringUtf8Coder.of()), parDoOperation.getBytesReadCoder()); } static ParallelInstruction createPartialGroupByKeyInstruction( diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java index 0251b075becf..dbdea61743dc 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java @@ -43,6 +43,7 @@ import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; +import org.apache.beam.runners.core.ElementByteSizeObservable; import org.apache.beam.runners.core.SideInputReader; import org.apache.beam.runners.core.StateInternals; import org.apache.beam.runners.core.StateNamespaceForTest; @@ -72,6 +73,12 @@ import org.apache.beam.runners.dataflow.worker.streaming.harness.StreamingCounters; import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputStateFetcherFactory; import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor; +import org.apache.beam.runners.dataflow.worker.util.common.worker.NativeReader; +import org.apache.beam.runners.dataflow.worker.util.common.worker.OutputObjectAndByteCounter; +import org.apache.beam.runners.dataflow.worker.util.common.worker.OutputReceiver; +import org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoFn; +import org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoOperation; +import org.apache.beam.runners.dataflow.worker.util.common.worker.ReadOperation; import org.apache.beam.runners.dataflow.worker.util.common.worker.WorkExecutor; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; import org.apache.beam.runners.dataflow.worker.windmill.client.getdata.FakeGetDataClient; @@ -943,4 +950,146 @@ public void testAdvance_stopsWhenCurrentWorkBatchingDisabled() throws Exception assertFalse(transitionListenerCalled.get()); verifyNoInteractions(mockExecutor); } + + @Test + public void testComputeSourceBytesProcessed_customSource() throws Exception { + CounterSet counterSet = new CounterSet(); + NameContext nameContext = NameContext.create("stage", "original", "system", "out"); + DataflowOperationContext opContext = executionContext.createOperationContext(nameContext); + + OutputReceiver readReceiver = new OutputReceiver(); + OutputObjectAndByteCounter customSourceCounter = + new OutputObjectAndByteCounter( + mock(ElementByteSizeObservable.class), counterSet, nameContext) + .countBytes("sourceBytesProcessCounterName"); + customSourceCounter.getByteCount().addValue(100L); + readReceiver.addOutputCounter("sourceBytesProcessCounterName", customSourceCounter); + + ReadOperation readOp = ReadOperation.forTest(mock(NativeReader.class), readReceiver, opContext); + + workExecutor = + IntrinsicMapTaskExecutor.forOperations( + Arrays.asList(readOp), mock(ExecutionStateTracker.class)); + + Work work = + createMockWork( + Windmill.WorkItem.newBuilder() + .setKey(ByteString.copyFromUtf8("key1")) + .setWorkToken(1L) + .build(), + Watermarks.builder().setInputDataWatermark(Instant.EPOCH).build()); + + start(work); + + long processed = executionContext.computeSourceBytesProcessed("sourceBytesProcessCounterName"); + assertEquals(100L, processed); + // Verified getAndReset resets to 0. + assertEquals(0L, executionContext.computeSourceBytesProcessed("sourceBytesProcessCounterName")); + } + + @Test + public void testComputeSourceBytesProcessed_splittableDoFn() throws Exception { + CounterSet counterSet = new CounterSet(); + NameContext nameContext = NameContext.create("stage", "original", "system", "out"); + DataflowOperationContext opContext = executionContext.createOperationContext(nameContext); + + OutputReceiver readReceiver = new OutputReceiver(); + ReadOperation readOp = ReadOperation.forTest(mock(NativeReader.class), readReceiver, opContext); + + ParDoFn mockParDoFn = mock(ParDoFn.class); + OutputReceiver parDoReceiver = new OutputReceiver(); + OutputObjectAndByteCounter sdfCounter = + new OutputObjectAndByteCounter( + mock(ElementByteSizeObservable.class), counterSet, nameContext) + .countBytes("sourceBytesProcessCounterName"); + sdfCounter.getByteCount().addValue(42L); + parDoReceiver.addOutputCounter("sourceBytesProcessCounterName", sdfCounter); + + ParDoOperation parDoOp = + new ParDoOperation( + mockParDoFn, new OutputReceiver[] {parDoReceiver}, opContext, mock(Coder.class)); + + workExecutor = + IntrinsicMapTaskExecutor.forOperations( + Arrays.asList(readOp, parDoOp), mock(ExecutionStateTracker.class)); + + Work work = + createMockWork( + Windmill.WorkItem.newBuilder() + .setKey(ByteString.copyFromUtf8("key1")) + .setWorkToken(1L) + .build(), + Watermarks.builder().setInputDataWatermark(Instant.EPOCH).build()); + + start(work); + + long processed = executionContext.computeSourceBytesProcessed("sourceBytesProcessCounterName"); + assertEquals(42L, processed); + // Verified getAndReset resets to 0. + assertEquals(0L, executionContext.computeSourceBytesProcessed("sourceBytesProcessCounterName")); + } + + @Test + public void testComputeSourceBytesProcessed_multipleParDoOperationsAndFlushState() + throws Exception { + NameContext nameContext = NameContext.create("stage", "original", "system", "out"); + DataflowOperationContext opContext = executionContext.createOperationContext(nameContext); + + OutputReceiver readReceiver = new OutputReceiver(); + ReadOperation readOp = ReadOperation.forTest(mock(NativeReader.class), readReceiver, opContext); + + ParDoFn mockParDoFn = mock(ParDoFn.class); + + OutputReceiver sdfReceiver1 = new OutputReceiver(); + OutputObjectAndByteCounter sdfCounter1 = + new OutputObjectAndByteCounter( + mock(ElementByteSizeObservable.class), new CounterSet(), nameContext) + .countBytes("sourceBytesProcessCounterName"); + sdfCounter1.getByteCount().addValue(40L); + sdfReceiver1.addOutputCounter("sourceBytesProcessCounterName", sdfCounter1); + ParDoOperation sdfParDoOp1 = + new ParDoOperation( + mockParDoFn, new OutputReceiver[] {sdfReceiver1}, opContext, mock(Coder.class)); + + OutputReceiver sdfReceiver2 = new OutputReceiver(); + OutputObjectAndByteCounter sdfCounter2 = + new OutputObjectAndByteCounter( + mock(ElementByteSizeObservable.class), new CounterSet(), nameContext) + .countBytes("sourceBytesProcessCounterName"); + sdfCounter2.getByteCount().addValue(60L); + sdfReceiver2.addOutputCounter("sourceBytesProcessCounterName", sdfCounter2); + ParDoOperation sdfParDoOp2 = + new ParDoOperation( + mockParDoFn, new OutputReceiver[] {sdfReceiver2}, opContext, mock(Coder.class)); + + OutputReceiver regularReceiver = new OutputReceiver(); + OutputObjectAndByteCounter regularCounter = + new OutputObjectAndByteCounter( + mock(ElementByteSizeObservable.class), new CounterSet(), nameContext) + .countBytes("sourceBytesProcessCounterName"); + regularCounter.getByteCount().addValue(999L); + regularReceiver.addOutputCounter("sourceBytesProcessCounterName", regularCounter); + ParDoOperation regularParDoOp = + new ParDoOperation(mockParDoFn, new OutputReceiver[] {regularReceiver}, opContext, null); + + workExecutor = + IntrinsicMapTaskExecutor.forOperations( + Arrays.asList(readOp, sdfParDoOp1, regularParDoOp, sdfParDoOp2), + mock(ExecutionStateTracker.class)); + + Work work = + createMockWork( + Windmill.WorkItem.newBuilder() + .setKey(ByteString.copyFromUtf8("key1")) + .setWorkToken(1L) + .build(), + Watermarks.builder().setInputDataWatermark(Instant.EPOCH).build()); + + start(work); + executionContext.finishKey(); + ExecuteWorkResult result = executionContext.flushStateAndReset(); + + assertEquals(1, result.workItemCommits().size()); + assertEquals(100L, result.workItemCommits().get(0).getSourceBytesProcessed()); + } } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutorTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutorTest.java index 40f4c7b0e715..0cfda97fd149 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutorTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutorTest.java @@ -62,6 +62,8 @@ import org.apache.beam.runners.dataflow.worker.counters.NameContext; import org.apache.beam.runners.dataflow.worker.profiler.ScopedProfiler.NoopProfileScope; import org.apache.beam.runners.dataflow.worker.util.common.worker.ExecutorTestUtils.TestReader; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.coders.VarIntCoder; import org.apache.beam.sdk.metrics.Metrics; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; @@ -573,4 +575,36 @@ public void testExceptionAndThenCloseAbortsJustOnce() throws Exception { // Closing after already closed should not call abort again. executor.close(); } + + @Test + public void testGetReadReportingParDoOperations() throws Exception { + TestOutputReceiver receiver = + new TestOutputReceiver(counterSet, NameContextsForTests.nameContextForTest()); + ReadOperation readOp = new TestReadOperation(receiver, createContext("ReadOperation")); + ParDoOperation regularParDo = + new ParDoOperation( + new NoopParDoFn(), new OutputReceiver[0], createContext("RegularParDo"), null); + ParDoOperation reportingParDo1 = + new ParDoOperation( + new NoopParDoFn(), + new OutputReceiver[0], + createContext("ReportingParDo1"), + StringUtf8Coder.of()); + ParDoOperation reportingParDo2 = + new ParDoOperation( + new NoopParDoFn(), + new OutputReceiver[0], + createContext("ReportingParDo2"), + VarIntCoder.of()); + + List operations = + Arrays.asList(readOp, regularParDo, reportingParDo1, reportingParDo2); + ExecutionStateTracker stateTracker = ExecutionStateTracker.newForTest(); + try (MapTaskExecutor executor = new MapTaskExecutor(operations, counterSet, stateTracker)) { + List reportingOps = executor.getReadReportingParDoOperations(); + assertThat(reportingOps, contains(reportingParDo1, reportingParDo2)); + assertEquals(StringUtf8Coder.of(), reportingOps.get(0).getBytesReadCoder()); + assertEquals(VarIntCoder.of(), reportingOps.get(1).getBytesReadCoder()); + } + } } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactoryTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactoryTest.java new file mode 100644 index 000000000000..ac20c4985db3 --- /dev/null +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactoryTest.java @@ -0,0 +1,323 @@ +/* + * 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.windmill.work.processing; + +import static org.apache.beam.runners.dataflow.util.Structs.addObject; +import static org.apache.beam.runners.dataflow.util.Structs.addString; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.mockito.Mockito.mock; + +import com.google.api.services.dataflow.model.InstructionInput; +import com.google.api.services.dataflow.model.InstructionOutput; +import com.google.api.services.dataflow.model.MapTask; +import com.google.api.services.dataflow.model.MultiOutputInfo; +import com.google.api.services.dataflow.model.ParDoInstruction; +import com.google.api.services.dataflow.model.ParallelInstruction; +import com.google.api.services.dataflow.model.ReadInstruction; +import com.google.api.services.dataflow.model.Source; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import org.apache.beam.runners.dataflow.options.DataflowWorkerHarnessOptions; +import org.apache.beam.runners.dataflow.util.CloudObject; +import org.apache.beam.runners.dataflow.util.CloudObjects; +import org.apache.beam.runners.dataflow.util.PropertyNames; +import org.apache.beam.runners.dataflow.worker.DataflowExecutionStateSampler; +import org.apache.beam.runners.dataflow.worker.DataflowMapTaskExecutor; +import org.apache.beam.runners.dataflow.worker.HotKeyLogger; +import org.apache.beam.runners.dataflow.worker.IntrinsicMapTaskExecutorFactory; +import org.apache.beam.runners.dataflow.worker.MultiKeyBundleOptions; +import org.apache.beam.runners.dataflow.worker.ReaderCache; +import org.apache.beam.runners.dataflow.worker.counters.Counter; +import org.apache.beam.runners.dataflow.worker.streaming.ComputationState; +import org.apache.beam.runners.dataflow.worker.streaming.ComputationWorkExecutor; +import org.apache.beam.runners.dataflow.worker.streaming.StageInfo; +import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingGlobalConfigHandleImpl; +import org.apache.beam.runners.dataflow.worker.streaming.harness.StreamingCounters; +import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputStateFetcherFactory; +import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor; +import org.apache.beam.runners.dataflow.worker.util.WorkerPropertyNames; +import org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoOperation; +import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillStateCache; +import org.apache.beam.runners.dataflow.worker.windmill.work.processing.failures.FailureTracker; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.coders.VarIntCoder; +import org.apache.beam.sdk.coders.VoidCoder; +import org.apache.beam.sdk.extensions.gcp.util.Transport; +import org.apache.beam.sdk.fn.IdGenerators; +import org.apache.beam.sdk.io.range.OffsetRange; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.DoFnSchemaInformation; +import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.util.DoFnInfo; +import org.apache.beam.sdk.util.SerializableUtils; +import org.apache.beam.sdk.util.StringUtils; +import org.apache.beam.sdk.values.TupleTag; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowedValues; +import org.apache.beam.sdk.values.WindowingStrategy; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for {@link ComputationWorkExecutorFactory}. */ +@RunWith(JUnit4.class) +public class ComputationWorkExecutorFactoryTest { + + private static class StringSplittableDoFn extends DoFn { + @ProcessElement + public void processElement(ProcessContext c, RestrictionTracker tracker) { + if (tracker.tryClaim(tracker.currentRestriction().getFrom())) { + c.output(c.element()); + } + } + + @GetInitialRestriction + public OffsetRange getInitialRestriction(@SuppressWarnings("unused") @Element String element) { + return new OffsetRange(0, 1); + } + } + + private static class IntegerSplittableDoFn extends DoFn { + @ProcessElement + public void processElement(ProcessContext c, RestrictionTracker tracker) { + if (tracker.tryClaim(tracker.currentRestriction().getFrom())) { + c.output(c.element().length()); + } + } + + @GetInitialRestriction + public OffsetRange getInitialRestriction(@SuppressWarnings("unused") @Element String element) { + return new OffsetRange(0, 1); + } + } + + private static class IdentityDoFn extends DoFn { + @ProcessElement + public void processElement(ProcessContext c) { + c.output(c.element()); + } + } + + private static ParallelInstruction createWindmillReadInstruction( + String name, Coder outputCoder) { + CloudObject encodedCoder = CloudObjects.asCloudObject(outputCoder, /* sdkComponents= */ null); + + Source cloudSource = new Source(); + cloudSource.setSpec(CloudObject.forClassName("UngroupedWindmillReader")); + cloudSource.setCodec(encodedCoder); + + ReadInstruction readInstruction = new ReadInstruction(); + readInstruction.setSource(cloudSource); + + InstructionOutput output = new InstructionOutput(); + output.setName(name + "_output"); + output.setCodec(encodedCoder); + output.setOriginalName(name + "Original"); + output.setSystemName(name + "System"); + + ParallelInstruction instruction = new ParallelInstruction(); + instruction.setName(name); + instruction.setSystemName(name + "System"); + instruction.setOriginalName(name + "Original"); + instruction.setRead(readInstruction); + instruction.setOutputs(Collections.singletonList(output)); + return instruction; + } + + private static ParallelInstruction createParDoInstruction( + int producerIndex, + int producerOutputNum, + String systemName, + DoFn fn, + Coder elementOutputCoder, + Coder> windowedOutputCoder, + boolean isSplittable) { + InstructionInput cloudInput = new InstructionInput(); + cloudInput.setProducerInstructionIndex(producerIndex); + cloudInput.setOutputNum(producerOutputNum); + + TupleTag mainOutputTag = new TupleTag<>(PropertyNames.OUTPUT); + String serializedFn = + StringUtils.byteArrayToJsonString( + SerializableUtils.serializeToByteArray( + DoFnInfo.forFn( + fn, + WindowingStrategy.globalDefault(), + null /* side input views */, + StringUtf8Coder.of() /* input coder */, + Collections.singletonMap(mainOutputTag, elementOutputCoder), + mainOutputTag, + DoFnSchemaInformation.create(), + Collections.emptyMap()))); + + CloudObject cloudUserFn = + CloudObject.forClassName(isSplittable ? "SplittableProcessFn" : "DoFn"); + addString(cloudUserFn, PropertyNames.SERIALIZED_FN, serializedFn); + if (isSplittable) { + addObject( + cloudUserFn, + WorkerPropertyNames.RESTRICTION_CODER, + CloudObjects.asCloudObject( + KvCoder.of(OffsetRange.Coder.of(), VoidCoder.of()), /* sdkComponents= */ null)); + } + + MultiOutputInfo multiOutputInfo = new MultiOutputInfo(); + multiOutputInfo.setTag(PropertyNames.OUTPUT); + + ParDoInstruction parDoInstruction = new ParDoInstruction(); + parDoInstruction.setInput(cloudInput); + parDoInstruction.setNumOutputs(1); + parDoInstruction.setMultiOutputInfos(ImmutableList.of(multiOutputInfo)); + parDoInstruction.setUserFn(cloudUserFn); + + InstructionOutput output = new InstructionOutput(); + output.setName(systemName + "_output"); + output.setCodec(CloudObjects.asCloudObject(windowedOutputCoder, /* sdkComponents= */ null)); + output.setOriginalName(systemName + "Original"); + output.setSystemName(systemName); + + ParallelInstruction instruction = new ParallelInstruction(); + instruction.setParDo(parDoInstruction); + instruction.setOutputs(Collections.singletonList(output)); + instruction.setName(systemName); + instruction.setSystemName(systemName); + instruction.setOriginalName(systemName + "Original"); + return instruction; + } + + @Test + public void + testCreateComputationWorkExecutor_multipleReadReportingParDoOperationsWithDifferentCoders() + throws Exception { + DataflowWorkerHarnessOptions options = + PipelineOptionsFactory.as(DataflowWorkerHarnessOptions.class); + ComputationWorkExecutorFactory factory = + new ComputationWorkExecutorFactory( + options, + IntrinsicMapTaskExecutorFactory.defaultFactory(), + mock(ReaderCache.class), + computationId -> mock(WindmillStateCache.ForComputation.class), + DataflowExecutionStateSampler.instance(), + StreamingCounters.create(), + mock(FailureTracker.class), + IdGenerators.decrementingLongs(), + new StreamingGlobalConfigHandleImpl(), + new HotKeyLogger(), + mock(SideInputStateFetcherFactory.class), + MultiKeyBundleOptions.EMPTY); + + Coder> windowedStringCoder = + WindowedValues.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE); + Coder> windowedIntCoder = + WindowedValues.getFullCoder(VarIntCoder.of(), GlobalWindow.Coder.INSTANCE); + + List instructions = + Arrays.asList( + createWindmillReadInstruction("Read", windowedStringCoder), + createParDoInstruction( + 0, + 0, + "SplittableParDoString", + new StringSplittableDoFn(), + StringUtf8Coder.of(), + windowedStringCoder, + true), + createParDoInstruction( + 0, + 0, + "SplittableParDoInt", + new IntegerSplittableDoFn(), + VarIntCoder.of(), + windowedIntCoder, + true), + createParDoInstruction( + 0, + 0, + "RegularParDo", + new IdentityDoFn(), + StringUtf8Coder.of(), + windowedStringCoder, + false)); + + MapTask mapTask = new MapTask(); + mapTask.setStageName("testStage"); + mapTask.setSystemName("testSystem"); + mapTask.setInstructions(instructions); + mapTask.setFactory(Transport.getJsonFactory()); + + StageInfo stageInfo = StageInfo.create("testStage", "testSystem"); + ComputationState computationState = + new ComputationState( + "computationId", + mapTask, + mock(BoundedQueueExecutor.class), + Collections.emptyMap(), + mock(WindmillStateCache.ForComputation.class)); + + ComputationWorkExecutor workExecutor = + factory.createComputationWorkExecutor(stageInfo, computationState, "workLatencyTrackingId"); + + DataflowMapTaskExecutor mapTaskExecutor = (DataflowMapTaskExecutor) workExecutor.workExecutor(); + List readReportingOps = mapTaskExecutor.getReadReportingParDoOperations(); + assertEquals(2, readReportingOps.size()); + + ParDoOperation stringSplittableOp = readReportingOps.get(0); + ParDoOperation intSplittableOp = readReportingOps.get(1); + assertEquals(windowedStringCoder, stringSplittableOp.getBytesReadCoder()); + assertEquals(windowedIntCoder, intSplittableOp.getBytesReadCoder()); + + String counterName = computationState.sourceBytesProcessCounterName(); + assertNotNull(stringSplittableOp.receivers[0].getOutputCounters().get(counterName)); + assertNotNull(intSplittableOp.receivers[0].getOutputCounters().get(counterName)); + + // Verify that the non-reporting ParDoOperation does not have the source bytes counter attached. + List allParDoOps = + ImmutableList.copyOf(Iterables.filter(mapTaskExecutor.operations, ParDoOperation.class)); + assertEquals(3, allParDoOps.size()); + ParDoOperation regularParDoOp = allParDoOps.get(2); + assertNull(regularParDoOp.getBytesReadCoder()); + assertNull(regularParDoOp.receivers[0].getOutputCounters().get(counterName)); + + // Process elements through each operation's receiver and verify each uses its own coder. + WindowedValue stringElem = WindowedValues.valueInGlobalWindow("hello"); + WindowedValue intElem = WindowedValues.valueInGlobalWindow(12345); + stringSplittableOp.receivers[0].process(stringElem); + intSplittableOp.receivers[0].process(intElem); + regularParDoOp.receivers[0].process(stringElem); + + long expectedStringBytes = CoderUtils.encodeToByteArray(windowedStringCoder, stringElem).length; + long expectedIntBytes = CoderUtils.encodeToByteArray(windowedIntCoder, intElem).length; + + @SuppressWarnings("unchecked") + Counter bytesReadCounter = + (Counter) mapTaskExecutor.getOutputCounters().getExistingCounter(counterName); + assertNotNull(bytesReadCounter); + assertEquals( + Long.valueOf(expectedStringBytes + expectedIntBytes), bytesReadCounter.getAggregate()); + } +} From f8cdf2a0a947d835addb370d6d73154a3420b5b3 Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Fri, 16 Oct 2026 13:23:14 -0700 Subject: [PATCH 2/9] Removes code not intended for this PR. Gets rid of duplicated trackAutoscalingParDoBytesRead method. --- ...oundedSplittableProcessElementInvoker.java | 2 +- ...edSplittableProcessElementInvokerTest.java | 70 ------------------- .../worker/SplittableProcessFnFactory.java | 4 +- .../worker/StreamingModeExecutionContext.java | 11 ++- .../ComputationWorkExecutorFactory.java | 60 +++++----------- 5 files changed, 31 insertions(+), 116 deletions(-) diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java index df1136a5d821..59d33c596c8a 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java @@ -182,7 +182,7 @@ public TimeDomain timeDomain(DoFn doFn) { @Override public OutputReceiver outputReceiver(DoFn doFn) { return DoFnOutputReceivers.windowedReceiver( - processContext, OutputBuilderSuppliers.supplierForElement(element), mainOutputTag); + processContext, OutputBuilderSuppliers.supplierForElement(element), null); } @Override diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvokerTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvokerTest.java index 32f21a32cb18..6ff40fae91dd 100644 --- a/runners/core-java/src/test/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvokerTest.java +++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvokerTest.java @@ -27,9 +27,7 @@ import static org.junit.Assert.assertNull; import static org.junit.Assert.assertTrue; -import java.util.ArrayList; import java.util.Collections; -import java.util.List; import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; import org.apache.beam.sdk.io.range.OffsetRange; @@ -42,7 +40,6 @@ import org.apache.beam.sdk.transforms.windowing.GlobalWindow; import org.apache.beam.sdk.transforms.windowing.PaneInfo; import org.apache.beam.sdk.util.WindowedValueMultiReceiver; -import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.sdk.values.WindowedValues; @@ -290,71 +287,4 @@ public void testBacklogBytesWhenDone() throws Exception { assertEquals(0.0, res.getBacklogBytes(), 0.001); assertNull(res.getResidualRestriction()); } - - private static class OutputReceiverFn extends DoFn { - @ProcessElement - public ProcessContinuation process( - OutputReceiver receiver, RestrictionTracker tracker) { - for (long i = tracker.currentRestriction().getFrom(); tracker.tryClaim(i); ++i) { - receiver.output(String.valueOf(i)); - } - return stop(); - } - - @GetInitialRestriction - public OffsetRange getInitialRestriction(@SuppressWarnings("unused") @Element Void element) { - return new OffsetRange(0, 3); - } - } - - @Test - public void testInvokeProcessElementWithOutputReceiverUsesMainOutputTag() throws Exception { - OutputReceiverFn fn = new OutputReceiverFn(); - TupleTag mainOutputTag = new TupleTag<>("mainOutputTag"); - List, String>> outputs = new ArrayList<>(); - - SplittableProcessElementInvoker invoker = - new OutputAndTimeBoundedSplittableProcessElementInvoker<>( - fn, - PipelineOptionsFactory.create(), - new WindowedValueMultiReceiver() { - @Override - public void output(TupleTag tag, WindowedValue output) { - outputs.add(KV.of(tag, (String) output.getValue())); - } - }, - mainOutputTag, - NullSideInputReader.empty(), - Executors.newSingleThreadScheduledExecutor(), - 1000, - Duration.standardMinutes(3), - () -> { - throw new UnsupportedOperationException("BundleFinalizer not configured for test."); - }); - - SplittableProcessElementInvoker.Result res = - invoker.invokeProcessElement( - DoFnInvokers.invokerFor(fn), - WindowedValues.of(null, Instant.now(), GlobalWindow.INSTANCE, PaneInfo.NO_FIRING), - new OffsetRangeTracker(new OffsetRange(0, 3)), - new WatermarkEstimator() { - @Override - public Instant currentWatermark() { - return GlobalWindow.TIMESTAMP_MIN_VALUE; - } - - @Override - public Void getState() { - return null; - } - }, - Collections.emptyMap()); - - assertFalse(res.getContinuation().shouldResume()); - assertNull(res.getResidualRestriction()); - assertEquals(3, outputs.size()); - assertEquals(KV.of(mainOutputTag, "0"), outputs.get(0)); - assertEquals(KV.of(mainOutputTag, "1"), outputs.get(1)); - assertEquals(KV.of(mainOutputTag, "2"), outputs.get(2)); - } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SplittableProcessFnFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SplittableProcessFnFactory.java index 42aaae3d95ba..b55d73cf5793 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SplittableProcessFnFactory.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SplittableProcessFnFactory.java @@ -151,11 +151,11 @@ public DoFnRunner>, OutputT> crea mainOutputTag, sideInputReader, ses.get(), - // Commit at least once every 5 seconds or 10k records. This keeps the watermark + // Commit at least once every 10 seconds or 10k records. This keeps the watermark // advancing smoothly, and ensures that not too much work will have to be reprocessed // in the event of a crash. 10000, - Duration.standardSeconds(5), + Duration.standardSeconds(10), stepContext::bundleFinalizer)); processFn.setBacklogBytesCallback(userStepContext::setBacklogBytes); DoFnRunner>, OutputT> simpleRunner = diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java index ada5dd1a4af1..d549d9163c7d 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java @@ -803,12 +803,17 @@ long computeSourceBytesProcessed(String sourceBytesCounterName) { if (!(workExecutor instanceof DataflowMapTaskExecutor)) { return 0L; } + // Note this is a problem for both custom source read operators and SDF ParDos: extracting + // the byte count counter values here races with StreamingWorkerStatusReporter, which also + // calls getAndReset on all modified counters to report to Dataflow. We should find a more + // dedicated counter that is only extracted here. DataflowMapTaskExecutor mapTaskExecutor = (DataflowMapTaskExecutor) workExecutor; long totalBytes = 0L; ReadOperation readOperation = mapTaskExecutor.getReadOperation(); if (readOperation != null && readOperation.receivers != null - && readOperation.receivers.length > 0) { + && readOperation.receivers.length > 0 + && readOperation.receivers[0] != null) { HashMap readOpCounters = readOperation.receivers[0].getOutputCounters(); totalBytes += @@ -817,7 +822,9 @@ long computeSourceBytesProcessed(String sourceBytesCounterName) { .orElse(0L); } for (ParDoOperation parDoOp : mapTaskExecutor.getReadReportingParDoOperations()) { - if (parDoOp.receivers != null && parDoOp.receivers.length > 0) { + if (parDoOp.receivers != null + && parDoOp.receivers.length > 0 + && parDoOp.receivers[0] != null) { HashMap splittableProcessOpCounters = parDoOp.receivers[0].getOutputCounters(); totalBytes += diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java index 622371d9fcaa..d70456f3d9a2 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java @@ -53,6 +53,7 @@ import org.apache.beam.runners.dataflow.worker.streaming.harness.StreamingCounters; import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputStateFetcherFactory; import org.apache.beam.runners.dataflow.worker.util.common.worker.MapTaskExecutor; +import org.apache.beam.runners.dataflow.worker.util.common.worker.Operation; import org.apache.beam.runners.dataflow.worker.util.common.worker.OutputObjectAndByteCounter; import org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoOperation; import org.apache.beam.runners.dataflow.worker.util.common.worker.ReadOperation; @@ -167,47 +168,27 @@ private static void trackAutoscalingBytesRead( MapTask mapTask, Nodes.ParallelInstructionNode readNode, Coder readCoder, - ReadOperation readOperation, + Operation operation, MapTaskExecutor mapTaskExecutor, String counterName) { - NameContext nameContext = + if (operation.receivers != null + && operation.receivers.length > 0 + && operation.receivers[0] != null) { + NameContext nameContext = NameContext.create( - mapTask.getStageName(), - readNode.getParallelInstruction().getOriginalName(), - readNode.getParallelInstruction().getSystemName(), - readNode.getParallelInstruction().getName()); - readOperation.receivers[0].addOutputCounter( + mapTask.getStageName(), + readNode.getParallelInstruction().getOriginalName(), + readNode.getParallelInstruction().getSystemName(), + readNode.getParallelInstruction().getName()); + operation.receivers[0].addOutputCounter( counterName, new OutputObjectAndByteCounter( - new IntrinsicMapTaskExecutorFactory.ElementByteSizeObservableCoder<>(readCoder), - mapTaskExecutor.getOutputCounters(), - nameContext) - .setSamplingPeriod(100) - .countBytes(counterName)); - } - - private static void trackAutoscalingParDoBytesRead( - MapTask mapTask, - Nodes.ParallelInstructionNode readNode, - Coder splittableProcessCoder, - ParDoOperation splittableParDoOp, - MapTaskExecutor mapTaskExecutor, - String counterName) { - NameContext nameContext = - NameContext.create( - mapTask.getStageName(), - readNode.getParallelInstruction().getOriginalName(), - readNode.getParallelInstruction().getSystemName(), - readNode.getParallelInstruction().getName()); - splittableParDoOp.receivers[0].addOutputCounter( - counterName, - new OutputObjectAndByteCounter( - new IntrinsicMapTaskExecutorFactory.ElementByteSizeObservableCoder<>( - splittableProcessCoder), - mapTaskExecutor.getOutputCounters(), - nameContext) - .setSamplingPeriod(100) - .countBytes(counterName)); + new IntrinsicMapTaskExecutorFactory.ElementByteSizeObservableCoder<>(readCoder), + mapTaskExecutor.getOutputCounters(), + nameContext) + .setSamplingPeriod(100) + .countBytes(counterName)); + } } private static ReadOperation getValidatedReadOperation(MapTaskExecutor mapTaskExecutor) { @@ -262,11 +243,8 @@ ComputationWorkExecutor createComputationWorkExecutor( } else { for (ParDoOperation parDoOp : mapTaskExecutor.getReadReportingParDoOperations()) { Coder bytesReadCoder = parDoOp.getBytesReadCoder(); - if (parDoOp.receivers != null - && parDoOp.receivers.length > 0 - && parDoOp.receivers[0] != null - && bytesReadCoder != null) { - trackAutoscalingParDoBytesRead( + if (bytesReadCoder != null) { + trackAutoscalingBytesRead( mapTask, readNode, bytesReadCoder, From 6acab498808f484ca15c693f4a337b35ac7391e0 Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Fri, 16 Oct 2026 13:48:56 -0700 Subject: [PATCH 3/9] Ran spotlessApply --- .../ComputationWorkExecutorFactory.java | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java index d70456f3d9a2..9356f36f02cc 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java @@ -175,19 +175,19 @@ private static void trackAutoscalingBytesRead( && operation.receivers.length > 0 && operation.receivers[0] != null) { NameContext nameContext = - NameContext.create( - mapTask.getStageName(), - readNode.getParallelInstruction().getOriginalName(), - readNode.getParallelInstruction().getSystemName(), - readNode.getParallelInstruction().getName()); + NameContext.create( + mapTask.getStageName(), + readNode.getParallelInstruction().getOriginalName(), + readNode.getParallelInstruction().getSystemName(), + readNode.getParallelInstruction().getName()); operation.receivers[0].addOutputCounter( - counterName, - new OutputObjectAndByteCounter( - new IntrinsicMapTaskExecutorFactory.ElementByteSizeObservableCoder<>(readCoder), - mapTaskExecutor.getOutputCounters(), - nameContext) - .setSamplingPeriod(100) - .countBytes(counterName)); + counterName, + new OutputObjectAndByteCounter( + new IntrinsicMapTaskExecutorFactory.ElementByteSizeObservableCoder<>(readCoder), + mapTaskExecutor.getOutputCounters(), + nameContext) + .setSamplingPeriod(100) + .countBytes(counterName)); } } From 45cde4e3a6d12c1803685ca98b8b93565afd913a Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Fri, 16 Oct 2026 14:06:46 -0700 Subject: [PATCH 4/9] Fixed failing test --- .../work/processing/ComputationWorkExecutorFactoryTest.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactoryTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactoryTest.java index ac20c4985db3..5ba10afc0d1b 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactoryTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactoryTest.java @@ -227,9 +227,9 @@ private static ParallelInstruction createParDoInstruction( mock(FailureTracker.class), IdGenerators.decrementingLongs(), new StreamingGlobalConfigHandleImpl(), - new HotKeyLogger(), + mock(HotKeyLogger.class), mock(SideInputStateFetcherFactory.class), - MultiKeyBundleOptions.EMPTY); + MultiKeyBundleOptions.fromOptions(options)); Coder> windowedStringCoder = WindowedValues.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE); From 293c6ac960dfb0ab806c2ed8534fa1788c03d73a Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Fri, 16 Oct 2026 14:35:27 -0700 Subject: [PATCH 5/9] Fixup setup for newly added tests. --- .../StreamingModeExecutionContextTest.java | 4 +++ .../ComputationWorkExecutorFactoryTest.java | 34 +++++++++++++------ 2 files changed, 28 insertions(+), 10 deletions(-) diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java index dbdea61743dc..99c598e56088 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java @@ -1085,6 +1085,10 @@ public void testComputeSourceBytesProcessed_multipleParDoOperationsAndFlushState .build(), Watermarks.builder().setInputDataWatermark(Instant.EPOCH).build()); + sdfParDoOp1.start(); + sdfParDoOp2.start(); + regularParDoOp.start(); + start(work); executionContext.finishKey(); ExecuteWorkResult result = executionContext.flushStateAndReset(); diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactoryTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactoryTest.java index 5ba10afc0d1b..69cbac36632e 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactoryTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactoryTest.java @@ -49,7 +49,8 @@ import org.apache.beam.runners.dataflow.worker.streaming.ComputationState; import org.apache.beam.runners.dataflow.worker.streaming.ComputationWorkExecutor; import org.apache.beam.runners.dataflow.worker.streaming.StageInfo; -import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingGlobalConfigHandleImpl; +import org.apache.beam.runners.dataflow.worker.streaming.config.FakeGlobalConfigHandle; +import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingGlobalConfig; import org.apache.beam.runners.dataflow.worker.streaming.harness.StreamingCounters; import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputStateFetcherFactory; import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor; @@ -226,9 +227,9 @@ private static ParallelInstruction createParDoInstruction( StreamingCounters.create(), mock(FailureTracker.class), IdGenerators.decrementingLongs(), - new StreamingGlobalConfigHandleImpl(), + new FakeGlobalConfigHandle(StreamingGlobalConfig.builder().build()), mock(HotKeyLogger.class), - mock(SideInputStateFetcherFactory.class), + SideInputStateFetcherFactory.fromOptions(options), MultiKeyBundleOptions.fromOptions(options)); Coder> windowedStringCoder = @@ -286,10 +287,17 @@ private static ParallelInstruction createParDoInstruction( List readReportingOps = mapTaskExecutor.getReadReportingParDoOperations(); assertEquals(2, readReportingOps.size()); - ParDoOperation stringSplittableOp = readReportingOps.get(0); - ParDoOperation intSplittableOp = readReportingOps.get(1); - assertEquals(windowedStringCoder, stringSplittableOp.getBytesReadCoder()); - assertEquals(windowedIntCoder, intSplittableOp.getBytesReadCoder()); + ParDoOperation stringSplittableOp = + readReportingOps.stream() + .filter(op -> windowedStringCoder.equals(op.getBytesReadCoder())) + .findFirst() + .orElseThrow( + () -> new AssertionError("Expected ParDoOperation with windowedStringCoder")); + ParDoOperation intSplittableOp = + readReportingOps.stream() + .filter(op -> windowedIntCoder.equals(op.getBytesReadCoder())) + .findFirst() + .orElseThrow(() -> new AssertionError("Expected ParDoOperation with windowedIntCoder")); String counterName = computationState.sourceBytesProcessCounterName(); assertNotNull(stringSplittableOp.receivers[0].getOutputCounters().get(counterName)); @@ -299,7 +307,11 @@ private static ParallelInstruction createParDoInstruction( List allParDoOps = ImmutableList.copyOf(Iterables.filter(mapTaskExecutor.operations, ParDoOperation.class)); assertEquals(3, allParDoOps.size()); - ParDoOperation regularParDoOp = allParDoOps.get(2); + ParDoOperation regularParDoOp = + allParDoOps.stream() + .filter(op -> !op.shouldReportBytesRead()) + .findFirst() + .orElseThrow(() -> new AssertionError("Expected non-reporting ParDoOperation")); assertNull(regularParDoOp.getBytesReadCoder()); assertNull(regularParDoOp.receivers[0].getOutputCounters().get(counterName)); @@ -310,8 +322,10 @@ private static ParallelInstruction createParDoInstruction( intSplittableOp.receivers[0].process(intElem); regularParDoOp.receivers[0].process(stringElem); - long expectedStringBytes = CoderUtils.encodeToByteArray(windowedStringCoder, stringElem).length; - long expectedIntBytes = CoderUtils.encodeToByteArray(windowedIntCoder, intElem).length; + long expectedStringBytes = + CoderUtils.encodeToByteArray(windowedStringCoder, stringElem, Coder.Context.NESTED).length; + long expectedIntBytes = + CoderUtils.encodeToByteArray(windowedIntCoder, intElem, Coder.Context.NESTED).length; @SuppressWarnings("unchecked") Counter bytesReadCounter = From 9f6dfe8ea86a4c7681c1e25d81fa355e15d01dff Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Wed, 21 Oct 2026 16:28:56 -0700 Subject: [PATCH 6/9] Add output counter to all outputs of source stages, not just the main output. --- .../ComputationWorkExecutorFactory.java | 26 +++++++++++-------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java index 9356f36f02cc..473772018999 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java @@ -55,6 +55,7 @@ import org.apache.beam.runners.dataflow.worker.util.common.worker.MapTaskExecutor; import org.apache.beam.runners.dataflow.worker.util.common.worker.Operation; import org.apache.beam.runners.dataflow.worker.util.common.worker.OutputObjectAndByteCounter; +import org.apache.beam.runners.dataflow.worker.util.common.worker.OutputReceiver; import org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoOperation; import org.apache.beam.runners.dataflow.worker.util.common.worker.ReadOperation; import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillStateCache; @@ -171,23 +172,26 @@ private static void trackAutoscalingBytesRead( Operation operation, MapTaskExecutor mapTaskExecutor, String counterName) { - if (operation.receivers != null - && operation.receivers.length > 0 - && operation.receivers[0] != null) { + if (operation.receivers != null && operation.receivers.length > 0) { NameContext nameContext = NameContext.create( mapTask.getStageName(), readNode.getParallelInstruction().getOriginalName(), readNode.getParallelInstruction().getSystemName(), readNode.getParallelInstruction().getName()); - operation.receivers[0].addOutputCounter( - counterName, - new OutputObjectAndByteCounter( - new IntrinsicMapTaskExecutorFactory.ElementByteSizeObservableCoder<>(readCoder), - mapTaskExecutor.getOutputCounters(), - nameContext) - .setSamplingPeriod(100) - .countBytes(counterName)); + for (OutputReceiver receiver : operation.receivers) { + if (receiver != null) { + receiver.addOutputCounter( + counterName, + new OutputObjectAndByteCounter( + new IntrinsicMapTaskExecutorFactory.ElementByteSizeObservableCoder<>( + readCoder), + mapTaskExecutor.getOutputCounters(), + nameContext) + .setSamplingPeriod(100) + .countBytes(counterName)); + } + } } } From 96aa8a0d2ecc1ad4e566c2c29f574db6c1851597 Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Tue, 27 Oct 2026 14:52:52 -0700 Subject: [PATCH 7/9] Changes how a SplittableFn is detected to not rely on class names. --- .../dataflow/worker/ForwardingParDoFn.java | 6 +++++ .../IntrinsicMapTaskExecutorFactory.java | 18 +++++++++++++- .../dataflow/worker/SimpleParDoFn.java | 24 ++++++------------- ...treamingKeyedWorkItemSideInputParDoFn.java | 24 ++++++------------- .../worker/StreamingModeExecutionContext.java | 5 ++-- .../util/common/worker/MapTaskExecutor.java | 11 ++++----- .../worker/util/common/worker/ParDoFn.java | 5 ++++ .../ComputationWorkExecutorFactory.java | 22 +++++++++-------- .../common/worker/MapTaskExecutorTest.java | 11 +++++---- .../ComputationWorkExecutorFactoryTest.java | 12 ++++------ 10 files changed, 71 insertions(+), 67 deletions(-) diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ForwardingParDoFn.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ForwardingParDoFn.java index 3a864d6caf2f..0c16325e56c4 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ForwardingParDoFn.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ForwardingParDoFn.java @@ -19,6 +19,7 @@ import org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoFn; import org.apache.beam.runners.dataflow.worker.util.common.worker.Receiver; +import org.apache.beam.sdk.util.DoFnInfo; import org.checkerframework.checker.nullness.qual.Nullable; /** @@ -62,4 +63,9 @@ public void finishBundle() throws Exception { public void abort() throws Exception { delegate.abort(); } + + @Override + public @Nullable DoFnInfo getDoFnInfo() throws Exception { + return delegate.getDoFnInfo(); + } } 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 7589a5036b58..1009cf8e0406 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 @@ -32,6 +32,7 @@ import java.util.List; import java.util.function.Function; import org.apache.beam.runners.core.ElementByteSizeObservable; +import org.apache.beam.runners.core.SplittableParDoViaKeyedWorkItems.ProcessFn; import org.apache.beam.runners.dataflow.util.CloudObject; import org.apache.beam.runners.dataflow.util.CloudObjects; import org.apache.beam.runners.dataflow.worker.counters.CounterFactory; @@ -64,6 +65,9 @@ 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.transforms.DoFn; +import org.apache.beam.sdk.transforms.reflect.DoFnSignatures; +import org.apache.beam.sdk.util.DoFnInfo; import org.apache.beam.sdk.util.common.ElementByteSizeObserver; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.WindowedValues.WindowedValueCoder; @@ -284,7 +288,7 @@ private OperationNode createParDoOperation( // An SDF's SplittableProcessFn behaves like a source, so we need to report its bytes read even // though it's not a read operator for autoscaling to work properly. Coder bytesReadCoder = null; - if ("SplittableProcessFn".equals(cloudUserFn.getClassName())) { + if (isSplittable(fn)) { Node firstOutput = Iterables.getFirst(network.successors(node), null); if (firstOutput instanceof OutputReceiverNode) { bytesReadCoder = ((OutputReceiverNode) firstOutput).getCoder(); @@ -295,6 +299,18 @@ private OperationNode createParDoOperation( new ParDoOperation(fn, receivers, operationContext, bytesReadCoder)); } + private static boolean isSplittable(ParDoFn fn) throws Exception { + DoFnInfo doFnInfo = fn.getDoFnInfo(); + if (doFnInfo == null) { + return false; + } + DoFn doFn = doFnInfo.getDoFn(); + if (doFn instanceof ProcessFn) { + doFn = ((ProcessFn) doFn).getFn(); + } + return DoFnSignatures.isSplittable(doFn); + } + private static TupleTag tupleTag(MultiOutputInfo multiOutputInfo) { return new TupleTag<>(multiOutputInfo.getTag()); } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFn.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFn.java index badc1533b2ad..7ced0ffb7451 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFn.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFn.java @@ -34,7 +34,6 @@ import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.sdk.values.WindowingStrategy; -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.Lists; import org.checkerframework.checker.nullness.qual.Nullable; @@ -112,7 +111,7 @@ protected void onStartKey() { // These elements are now processed. Register cleanup timers for all the unblocked // windows. helpers.registerStateCleanup( - (WindowingStrategy) getDoFnInfo().getWindowingStrategy(), + (WindowingStrategy) helpers.fnInfo.getWindowingStrategy(), (Collection) unblockedElement.getWindows()); } } @@ -163,7 +162,7 @@ protected void onProcessWindowedValue(WindowedValue elem) { } if (hasState) { helpers.registerStateCleanup( - (WindowingStrategy) getDoFnInfo().getWindowingStrategy(), windowsProcessed); + (WindowingStrategy) helpers.fnInfo.getWindowingStrategy(), windowsProcessed); } } @@ -176,11 +175,7 @@ public void processTimers() throws Exception { // windowCoder (allowing us to get it without a DoFnInfo instance) or (2) check whether timers // exist without actually decoding them. Coder windowCoder = - (Coder) - (helpers.fnInfo != null ? helpers.fnInfo : helpers.doFnInstanceManager.peek()) - .getWindowingStrategy() - .getWindowFn() - .windowCoder(); + (Coder) getDoFnInfo().getWindowingStrategy().getWindowFn().windowCoder(); helpers.processTimers( SimpleParDoFnHelpers.TimerType.USER, helpers.userStepContext, @@ -209,14 +204,9 @@ public void abort() throws Exception { helpers.abort(); } - /** - * Returns the {@link DoFnInfo} currently being used by this {@link SimpleParDoFn}. - * - *

May be null if no element has been processed yet, or if the {@link SimpleParDoFn} has - * finished. - */ - @VisibleForTesting - @Nullable DoFnInfo getDoFnInfo() { - return helpers.fnInfo; + /** Returns the {@link DoFnInfo} used by this {@link SimpleParDoFn}. */ + @Override + public DoFnInfo getDoFnInfo() throws Exception { + return helpers.fnInfo != null ? helpers.fnInfo : helpers.doFnInstanceManager.peek(); } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputParDoFn.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputParDoFn.java index 5e61ee3a57f3..882ee7a8cacc 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputParDoFn.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputParDoFn.java @@ -45,7 +45,6 @@ import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.sdk.values.WindowingStrategy; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; import org.checkerframework.checker.nullness.qual.Nullable; @SuppressWarnings({ @@ -138,7 +137,7 @@ protected void onStartKey() { // These elements are now processed. Register cleanup timers for all the unblocked // windows. helpers.registerStateCleanup( - (WindowingStrategy) getDoFnInfo().getWindowingStrategy(), windows); + (WindowingStrategy) helpers.fnInfo.getWindowingStrategy(), windows); } }); } @@ -174,11 +173,7 @@ public void processTimers() throws Exception { // windowCoder (allowing us to get it without a DoFnInfo instance) or (2) check whether timers // exist without actually decoding them. Coder windowCoder = - (Coder) - (helpers.fnInfo != null ? helpers.fnInfo : helpers.doFnInstanceManager.peek()) - .getWindowingStrategy() - .getWindowFn() - .windowCoder(); + (Coder) getDoFnInfo().getWindowingStrategy().getWindowFn().windowCoder(); helpers.processTimers( SimpleParDoFnHelpers.TimerType.FAIL_USER, helpers.userStepContext, @@ -231,18 +226,13 @@ protected void onProcessWindowedValue(WindowedValue> el } if (hasState) { helpers.registerStateCleanup( - (WindowingStrategy) getDoFnInfo().getWindowingStrategy(), windowsProcessed); + (WindowingStrategy) helpers.fnInfo.getWindowingStrategy(), windowsProcessed); } } - /** - * Returns the {@link DoFnInfo} currently being used by this {@link SimpleParDoFn}. - * - *

May be null if no element has been processed yet, or if the {@link SimpleParDoFn} has - * finished. - */ - @VisibleForTesting - @Nullable DoFnInfo getDoFnInfo() { - return helpers.fnInfo; + /** Returns the {@link DoFnInfo} used by this {@link StreamingKeyedWorkItemSideInputParDoFn}. */ + @Override + public DoFnInfo getDoFnInfo() throws Exception { + return helpers.fnInfo != null ? helpers.fnInfo : helpers.doFnInstanceManager.peek(); } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java index d549d9163c7d..3c37aae7306c 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java @@ -821,8 +821,9 @@ long computeSourceBytesProcessed(String sourceBytesCounterName) { .map(counter -> ((OutputObjectAndByteCounter) counter).getByteCount().getAndReset()) .orElse(0L); } - for (ParDoOperation parDoOp : mapTaskExecutor.getReadReportingParDoOperations()) { - if (parDoOp.receivers != null + for (ParDoOperation parDoOp : mapTaskExecutor.getParDoOperations()) { + if (parDoOp.shouldReportBytesRead() + && parDoOp.receivers != null && parDoOp.receivers.length > 0 && parDoOp.receivers[0] != null) { HashMap splittableProcessOpCounters = diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutor.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutor.java index 68c94dc89f3c..919b0acf4de3 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutor.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutor.java @@ -145,19 +145,16 @@ public ReadOperation getReadOperation() { return (ReadOperation) readOperation; } - public List getReadReportingParDoOperations() { - List reportingParDoOperations = Lists.newArrayList(); + public List getParDoOperations() { + List parDoOperations = Lists.newArrayList(); if (operations != null) { for (Operation op : operations) { if (op instanceof ParDoOperation) { - ParDoOperation parDoOp = (ParDoOperation) op; - if (parDoOp.shouldReportBytesRead()) { - reportingParDoOperations.add(parDoOp); - } + parDoOperations.add((ParDoOperation) op); } } } - return reportingParDoOperations; + return parDoOperations; } public boolean supportsRestart() { diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoFn.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoFn.java index 75ecdb67dd42..86473ed0a337 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoFn.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoFn.java @@ -17,6 +17,7 @@ */ package org.apache.beam.runners.dataflow.worker.util.common.worker; +import org.apache.beam.sdk.util.DoFnInfo; import org.checkerframework.checker.nullness.qual.Nullable; /** @@ -37,4 +38,8 @@ public interface ParDoFn { void finishBundle() throws Exception; void abort() throws Exception; + + default @Nullable DoFnInfo getDoFnInfo() throws Exception { + return null; + } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java index 473772018999..8b3541f9b9c3 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java @@ -245,16 +245,18 @@ ComputationWorkExecutor createComputationWorkExecutor( mapTaskExecutor, computationState.sourceBytesProcessCounterName()); } else { - for (ParDoOperation parDoOp : mapTaskExecutor.getReadReportingParDoOperations()) { - Coder bytesReadCoder = parDoOp.getBytesReadCoder(); - if (bytesReadCoder != null) { - trackAutoscalingBytesRead( - mapTask, - readNode, - bytesReadCoder, - parDoOp, - mapTaskExecutor, - computationState.sourceBytesProcessCounterName()); + for (ParDoOperation parDoOp : mapTaskExecutor.getParDoOperations()) { + if (parDoOp.shouldReportBytesRead()) { + Coder bytesReadCoder = parDoOp.getBytesReadCoder(); + if (bytesReadCoder != null) { + trackAutoscalingBytesRead( + mapTask, + readNode, + bytesReadCoder, + parDoOp, + mapTaskExecutor, + computationState.sourceBytesProcessCounterName()); + } } } } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutorTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutorTest.java index 0cfda97fd149..f707c0dd194f 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutorTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutorTest.java @@ -577,7 +577,7 @@ public void testExceptionAndThenCloseAbortsJustOnce() throws Exception { } @Test - public void testGetReadReportingParDoOperations() throws Exception { + public void testGetParDoOperations() throws Exception { TestOutputReceiver receiver = new TestOutputReceiver(counterSet, NameContextsForTests.nameContextForTest()); ReadOperation readOp = new TestReadOperation(receiver, createContext("ReadOperation")); @@ -601,10 +601,11 @@ public void testGetReadReportingParDoOperations() throws Exception { Arrays.asList(readOp, regularParDo, reportingParDo1, reportingParDo2); ExecutionStateTracker stateTracker = ExecutionStateTracker.newForTest(); try (MapTaskExecutor executor = new MapTaskExecutor(operations, counterSet, stateTracker)) { - List reportingOps = executor.getReadReportingParDoOperations(); - assertThat(reportingOps, contains(reportingParDo1, reportingParDo2)); - assertEquals(StringUtf8Coder.of(), reportingOps.get(0).getBytesReadCoder()); - assertEquals(VarIntCoder.of(), reportingOps.get(1).getBytesReadCoder()); + List parDoOps = executor.getParDoOperations(); + assertThat(parDoOps, contains(regularParDo, reportingParDo1, reportingParDo2)); + Assert.assertNull(parDoOps.get(0).getBytesReadCoder()); + assertEquals(StringUtf8Coder.of(), parDoOps.get(1).getBytesReadCoder()); + assertEquals(VarIntCoder.of(), parDoOps.get(2).getBytesReadCoder()); } } } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactoryTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactoryTest.java index 69cbac36632e..5c0adf1e7cbd 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactoryTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactoryTest.java @@ -80,7 +80,6 @@ import org.apache.beam.sdk.values.WindowedValues; import org.apache.beam.sdk.values.WindowingStrategy; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -284,17 +283,17 @@ private static ParallelInstruction createParDoInstruction( factory.createComputationWorkExecutor(stageInfo, computationState, "workLatencyTrackingId"); DataflowMapTaskExecutor mapTaskExecutor = (DataflowMapTaskExecutor) workExecutor.workExecutor(); - List readReportingOps = mapTaskExecutor.getReadReportingParDoOperations(); - assertEquals(2, readReportingOps.size()); + List allParDoOps = mapTaskExecutor.getParDoOperations(); + assertEquals(3, allParDoOps.size()); ParDoOperation stringSplittableOp = - readReportingOps.stream() + allParDoOps.stream() .filter(op -> windowedStringCoder.equals(op.getBytesReadCoder())) .findFirst() .orElseThrow( () -> new AssertionError("Expected ParDoOperation with windowedStringCoder")); ParDoOperation intSplittableOp = - readReportingOps.stream() + allParDoOps.stream() .filter(op -> windowedIntCoder.equals(op.getBytesReadCoder())) .findFirst() .orElseThrow(() -> new AssertionError("Expected ParDoOperation with windowedIntCoder")); @@ -304,9 +303,6 @@ private static ParallelInstruction createParDoInstruction( assertNotNull(intSplittableOp.receivers[0].getOutputCounters().get(counterName)); // Verify that the non-reporting ParDoOperation does not have the source bytes counter attached. - List allParDoOps = - ImmutableList.copyOf(Iterables.filter(mapTaskExecutor.operations, ParDoOperation.class)); - assertEquals(3, allParDoOps.size()); ParDoOperation regularParDoOp = allParDoOps.stream() .filter(op -> !op.shouldReportBytesRead()) From 1e4db1f3f7bd98c83e77538450eca55e1cb78931 Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Wed, 28 Oct 2026 11:56:02 -0700 Subject: [PATCH 8/9] Addressing PR comments. Removes default implementation of getDoFnInfo and renaming helpers.doFnInfo to activeDoFnInfo. --- .../worker/AssignWindowsParDoFnFactory.java | 6 +++ .../worker/BatchModeUngroupingParDoFn.java | 7 ++++ ...reateIsmShardKeyAndSortKeyDoFnFactory.java | 7 ++++ .../worker/GroupAlsoByWindowsParDoFn.java | 6 +++ .../PairWithConstantKeyDoFnFactory.java | 7 ++++ .../worker/PartialGroupByKeyParDoFns.java | 11 +++++ ...eifyTimestampAndWindowsParDoFnFactory.java | 6 +++ .../dataflow/worker/SimpleParDoFn.java | 40 ++++++++++++++----- .../dataflow/worker/SimpleParDoFnHelpers.java | 31 +++++++------- ...treamingKeyedWorkItemSideInputParDoFn.java | 38 +++++++++++++----- ...StreamingPCollectionViewWriterParDoFn.java | 7 ++++ .../ToIsmRecordForMultimapDoFnFactory.java | 7 ++++ .../dataflow/worker/ValuesDoFnFactory.java | 7 ++++ .../worker/util/common/worker/ParDoFn.java | 8 ++-- .../SimplePartialGroupByKeyParDoFn.java | 6 +++ .../worker/IntrinsicMapTaskExecutorTest.java | 7 ++++ .../worker/UserParDoFnFactoryTest.java | 13 +++--- .../common/worker/MapTaskExecutorTest.java | 7 ++++ .../common/worker/ParDoOperationTest.java | 7 ++++ 19 files changed, 183 insertions(+), 45 deletions(-) diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/AssignWindowsParDoFnFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/AssignWindowsParDoFnFactory.java index 673130a6048a..897d707456e2 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/AssignWindowsParDoFnFactory.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/AssignWindowsParDoFnFactory.java @@ -31,6 +31,7 @@ import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.transforms.windowing.WindowFn; +import org.apache.beam.sdk.util.DoFnInfo; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.sdk.values.WindowedValues; @@ -131,5 +132,10 @@ public void finishBundle() throws Exception { public void abort() throws Exception { receiver = null; } + + @Override + public @Nullable DoFnInfo getDoFnInfo() { + return null; + } } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/BatchModeUngroupingParDoFn.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/BatchModeUngroupingParDoFn.java index c3cf6d9e67e2..df1438b36e00 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/BatchModeUngroupingParDoFn.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/BatchModeUngroupingParDoFn.java @@ -19,8 +19,10 @@ import org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoFn; import org.apache.beam.runners.dataflow.worker.util.common.worker.Receiver; +import org.apache.beam.sdk.util.DoFnInfo; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.WindowedValue; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; /** @@ -85,4 +87,9 @@ public void finishBundle() throws Exception { public void abort() throws Exception { underlyingParDoFn.abort(); } + + @Override + public @Nullable DoFnInfo getDoFnInfo() throws Exception { + return underlyingParDoFn.getDoFnInfo(); + } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/CreateIsmShardKeyAndSortKeyDoFnFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/CreateIsmShardKeyAndSortKeyDoFnFactory.java index afe2e28b2a67..9d28b4720aa0 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/CreateIsmShardKeyAndSortKeyDoFnFactory.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/CreateIsmShardKeyAndSortKeyDoFnFactory.java @@ -32,10 +32,12 @@ import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.DoFnInfo; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.checkerframework.checker.nullness.qual.Nullable; /** * A {@link ParDoFnFactory} that creates a system {@link ParDoFn} responsible for limiting the users @@ -122,5 +124,10 @@ public void finishBundle() {} @Override public void abort() {} + + @Override + public @Nullable DoFnInfo getDoFnInfo() { + return null; + } } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/GroupAlsoByWindowsParDoFn.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/GroupAlsoByWindowsParDoFn.java index 6c92b7f10ea1..f4364f006c51 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/GroupAlsoByWindowsParDoFn.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/GroupAlsoByWindowsParDoFn.java @@ -34,6 +34,7 @@ import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.StreamingOptions; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.DoFnInfo; import org.apache.beam.sdk.util.WindowedValueReceiver; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollectionView; @@ -160,6 +161,11 @@ public void abort() throws Exception { fnRunner = null; } + @Override + public @Nullable DoFnInfo getDoFnInfo() { + return null; + } + /** * Composes and returns a {@link DoFnRunner} based on the parameters. * diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PairWithConstantKeyDoFnFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PairWithConstantKeyDoFnFactory.java index 425184a4a126..5a0406d9cdd7 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PairWithConstantKeyDoFnFactory.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PairWithConstantKeyDoFnFactory.java @@ -32,9 +32,11 @@ import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.util.DoFnInfo; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.WindowedValue; +import org.checkerframework.checker.nullness.qual.Nullable; /** * A {@link ParDoFnFactory} which returns a {@link ParDoFn} that transforms all {@code @@ -106,5 +108,10 @@ public void finishBundle() {} @Override public void abort() {} + + @Override + public @Nullable DoFnInfo getDoFnInfo() { + return null; + } } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PartialGroupByKeyParDoFns.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PartialGroupByKeyParDoFns.java index a6d7810412a1..c81ad13ef38b 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PartialGroupByKeyParDoFns.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PartialGroupByKeyParDoFns.java @@ -43,6 +43,7 @@ import org.apache.beam.sdk.state.BagState; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.util.AppliedCombineFn; +import org.apache.beam.sdk.util.DoFnInfo; import org.apache.beam.sdk.util.SerializableUtils; import org.apache.beam.sdk.util.common.ElementByteSizeObserver; import org.apache.beam.sdk.values.KV; @@ -327,6 +328,11 @@ public void finishBundle() throws Exception { @Override public void abort() throws Exception {} + + @Override + public @Nullable DoFnInfo getDoFnInfo() { + return null; + } } static class StreamingSideInputPGBKParDoFn @@ -391,5 +397,10 @@ public void finishBundle() throws Exception { @Override public void abort() throws Exception {} + + @Override + public @Nullable DoFnInfo getDoFnInfo() { + return null; + } } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ReifyTimestampAndWindowsParDoFnFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ReifyTimestampAndWindowsParDoFnFactory.java index 0438b525b6b9..eb3d6fcdfeab 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ReifyTimestampAndWindowsParDoFnFactory.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ReifyTimestampAndWindowsParDoFnFactory.java @@ -26,6 +26,7 @@ import org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoFn; import org.apache.beam.runners.dataflow.worker.util.common.worker.Receiver; import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.util.DoFnInfo; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.WindowedValue; @@ -98,5 +99,10 @@ public void finishBundle() throws Exception { public void abort() throws Exception { this.receiver = null; } + + @Override + public @Nullable DoFnInfo getDoFnInfo() { + return null; + } } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFn.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFn.java index 7ced0ffb7451..4b43a84174f5 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFn.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFn.java @@ -34,6 +34,7 @@ import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.sdk.values.WindowingStrategy; +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.Lists; import org.checkerframework.checker.nullness.qual.Nullable; @@ -96,9 +97,9 @@ protected void onStartKey() { sideInputProcessor = new StreamingSideInputProcessor<>( new StreamingSideInputFetcher( - helpers.fnInfo.getSideInputViews(), - helpers.fnInfo.getInputCoder(), - (WindowingStrategy) helpers.fnInfo.getWindowingStrategy(), + helpers.activeFnInfo.getSideInputViews(), + helpers.activeFnInfo.getInputCoder(), + (WindowingStrategy) helpers.activeFnInfo.getWindowingStrategy(), (StreamingModeExecutionContext.StreamingModeStepContext) helpers.userStepContext)); @@ -111,7 +112,7 @@ protected void onStartKey() { // These elements are now processed. Register cleanup timers for all the unblocked // windows. helpers.registerStateCleanup( - (WindowingStrategy) helpers.fnInfo.getWindowingStrategy(), + (WindowingStrategy) helpers.activeFnInfo.getWindowingStrategy(), (Collection) unblockedElement.getWindows()); } } @@ -162,7 +163,7 @@ protected void onProcessWindowedValue(WindowedValue elem) { } if (hasState) { helpers.registerStateCleanup( - (WindowingStrategy) helpers.fnInfo.getWindowingStrategy(), windowsProcessed); + (WindowingStrategy) helpers.activeFnInfo.getWindowingStrategy(), windowsProcessed); } } @@ -170,12 +171,18 @@ protected void onProcessWindowedValue(WindowedValue elem) { public void processTimers() throws Exception { // Note: We need to get windowCoder to decode the timers. If we haven't already deserialized - // the fnInfo, we peek at a new instance to retrieve that. If this extra deserialization becomes - // excessively costly, we could either (1) have the DoFnInstanceManager remember the associated - // windowCoder (allowing us to get it without a DoFnInfo instance) or (2) check whether timers - // exist without actually decoding them. + // the activeFnInfo, we peek at a new instance to retrieve that. If this extra deserialization + // becomes excessively costly, we could either (1) have the DoFnInstanceManager remember the + // associated windowCoder (allowing us to get it without a DoFnInfo instance) or (2) check + // whether timers exist without actually decoding them. Coder windowCoder = - (Coder) getDoFnInfo().getWindowingStrategy().getWindowFn().windowCoder(); + (Coder) + (helpers.activeFnInfo != null + ? helpers.activeFnInfo + : helpers.doFnInstanceManager.peek()) + .getWindowingStrategy() + .getWindowFn() + .windowCoder(); helpers.processTimers( SimpleParDoFnHelpers.TimerType.USER, helpers.userStepContext, @@ -207,6 +214,17 @@ public void abort() throws Exception { /** Returns the {@link DoFnInfo} used by this {@link SimpleParDoFn}. */ @Override public DoFnInfo getDoFnInfo() throws Exception { - return helpers.fnInfo != null ? helpers.fnInfo : helpers.doFnInstanceManager.peek(); + return helpers.doFnInstanceManager.peek(); + } + + /** + * Returns the {@link DoFnInfo} currently being used by this {@link SimpleParDoFn}. + * + *

May be null if no element has been processed yet, or if the {@link SimpleParDoFn} has + * finished. + */ + @VisibleForTesting + @Nullable DoFnInfo getActiveDoFnInfo() { + return helpers.activeFnInfo; } } 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 15bfba9bbc49..a69b3af7e256 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 @@ -97,7 +97,7 @@ class SimpleParDoFnHelpers { // Various DoFn helpers, null between bundles @Nullable DoFnRunner fnRunner; - @Nullable DoFnInfo fnInfo; + @Nullable DoFnInfo activeFnInfo; private Receiver @Nullable [] receivers; // This may additionally be null if it is not a real DoFn but an OldDoFn or @@ -212,20 +212,20 @@ public void output(TupleTag tag, WindowedValue output) { } } }; - fnInfo = (DoFnInfo) doFnInstanceManager.get(); - fnSignature = DoFnSignatures.getSignature(fnInfo.getDoFn().getClass()); + activeFnInfo = (DoFnInfo) doFnInstanceManager.get(); + fnSignature = DoFnSignatures.getSignature(activeFnInfo.getDoFn().getClass()); fnRunner = runnerFactory.createRunner( - fnInfo.getDoFn(), + activeFnInfo.getDoFn(), options, mainOutputTag, sideOutputTags, - fnInfo.getSideInputViews(), + activeFnInfo.getSideInputViews(), sideInputReader, - fnInfo.getInputCoder(), - fnInfo.getOutputCoders(), - fnInfo.getWindowingStrategy(), + activeFnInfo.getInputCoder(), + activeFnInfo.getOutputCoders(), + activeFnInfo.getWindowingStrategy(), stepContext, userStepContext, outputManager, @@ -240,18 +240,18 @@ void finishBundle(StreamingSideInputProcessor sideInputProcessor) throws E if (sideInputProcessor != null) { sideInputProcessor.handleFinishBundle(); } - doFnInstanceManager.complete(fnInfo); + doFnInstanceManager.complete(activeFnInfo); fnRunner = null; - fnInfo = null; + activeFnInfo = null; fnSignature = null; sideInputProcessor = null; } } void abort() throws Exception { - doFnInstanceManager.abort(fnInfo); + doFnInstanceManager.abort(activeFnInfo); fnRunner = null; - fnInfo = null; + activeFnInfo = null; } @VisibleForTesting static final String CLEANUP_TIMER_ID = "cleanup-timer"; @@ -366,7 +366,8 @@ private void processSystemTimer( } BoundedWindow window = ((StateNamespaces.WindowNamespace) timer.getNamespace()).getWindow(); - Instant targetTime = earliestAllowableCleanupTime(window, fnInfo.getWindowingStrategy()); + Instant targetTime = + earliestAllowableCleanupTime(window, activeFnInfo.getWindowingStrategy()); checkState( !targetTime.isAfter(timer.getTimestamp()), @@ -385,12 +386,12 @@ private void processSystemTimer( try { tag = StateTags.tagForSpec( - stateDecl.id(), (StateSpec) stateDecl.field().get(fnInfo.getDoFn())); + stateDecl.id(), (StateSpec) stateDecl.field().get(activeFnInfo.getDoFn())); } catch (IllegalAccessException e) { throw new RuntimeException( String.format( "Error accessing %s for %s", - StateSpec.class.getName(), fnInfo.getDoFn().getClass().getName()), + StateSpec.class.getName(), activeFnInfo.getDoFn().getClass().getName()), e); } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputParDoFn.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputParDoFn.java index 882ee7a8cacc..3f9b695f89a5 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputParDoFn.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputParDoFn.java @@ -45,6 +45,7 @@ import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.sdk.values.WindowingStrategy; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; import org.checkerframework.checker.nullness.qual.Nullable; @SuppressWarnings({ @@ -108,9 +109,9 @@ protected void onStartKey() { sideInputProcessor = new StreamingSideInputProcessor<>( new StreamingSideInputFetcher( - helpers.fnInfo.getSideInputViews(), + helpers.activeFnInfo.getSideInputViews(), inputCoder, - (WindowingStrategy) helpers.fnInfo.getWindowingStrategy(), + (WindowingStrategy) helpers.activeFnInfo.getWindowingStrategy(), (StreamingModeExecutionContext.StreamingModeStepContext) helpers.userStepContext)); } @@ -137,7 +138,7 @@ protected void onStartKey() { // These elements are now processed. Register cleanup timers for all the unblocked // windows. helpers.registerStateCleanup( - (WindowingStrategy) helpers.fnInfo.getWindowingStrategy(), windows); + (WindowingStrategy) helpers.activeFnInfo.getWindowingStrategy(), windows); } }); } @@ -168,12 +169,18 @@ public void processElement(Object untypedElem) throws Exception { public void processTimers() throws Exception { // Note: We need to get windowCoder to decode the timers. If we haven't already deserialized - // the fnInfo, we peek at a new instance to retrieve that. If this extra deserialization becomes - // excessively costly, we could either (1) have the DoFnInstanceManager remember the associated - // windowCoder (allowing us to get it without a DoFnInfo instance) or (2) check whether timers - // exist without actually decoding them. + // the activeFnInfo, we peek at a new instance to retrieve that. If this extra deserialization + // becomes excessively costly, we could either (1) have the DoFnInstanceManager remember the + // associated windowCoder (allowing us to get it without a DoFnInfo instance) or (2) check + // whether timers exist without actually decoding them. Coder windowCoder = - (Coder) getDoFnInfo().getWindowingStrategy().getWindowFn().windowCoder(); + (Coder) + (helpers.activeFnInfo != null + ? helpers.activeFnInfo + : helpers.doFnInstanceManager.peek()) + .getWindowingStrategy() + .getWindowFn() + .windowCoder(); helpers.processTimers( SimpleParDoFnHelpers.TimerType.FAIL_USER, helpers.userStepContext, @@ -226,13 +233,24 @@ protected void onProcessWindowedValue(WindowedValue> el } if (hasState) { helpers.registerStateCleanup( - (WindowingStrategy) helpers.fnInfo.getWindowingStrategy(), windowsProcessed); + (WindowingStrategy) helpers.activeFnInfo.getWindowingStrategy(), windowsProcessed); } } /** Returns the {@link DoFnInfo} used by this {@link StreamingKeyedWorkItemSideInputParDoFn}. */ @Override public DoFnInfo getDoFnInfo() throws Exception { - return helpers.fnInfo != null ? helpers.fnInfo : helpers.doFnInstanceManager.peek(); + return helpers.doFnInstanceManager.peek(); + } + + /** + * Returns the {@link DoFnInfo} currently being used by this {@link SimpleParDoFn}. + * + *

May be null if no element has been processed yet, or if the {@link SimpleParDoFn} has + * finished. + */ + @VisibleForTesting + @Nullable DoFnInfo getActiveDoFnInfo() { + return helpers.activeFnInfo; } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingPCollectionViewWriterParDoFn.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingPCollectionViewWriterParDoFn.java index 6b51427bb930..21ce2574bc81 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingPCollectionViewWriterParDoFn.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingPCollectionViewWriterParDoFn.java @@ -24,9 +24,11 @@ import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.coders.IterableCoder; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.DoFnInfo; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; +import org.checkerframework.checker.nullness.qual.Nullable; /** * A {@link ParDoFn} that writes side input data using {@link @@ -83,4 +85,9 @@ public void finishBundle() throws Exception {} public void abort() throws Exception { // no-op as long as it remains stateless } + + @Override + public @Nullable DoFnInfo getDoFnInfo() { + return null; + } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ToIsmRecordForMultimapDoFnFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ToIsmRecordForMultimapDoFnFactory.java index 261ed69fb5d7..c275fb8e3d09 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ToIsmRecordForMultimapDoFnFactory.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ToIsmRecordForMultimapDoFnFactory.java @@ -35,10 +35,12 @@ import org.apache.beam.sdk.coders.KvCoder; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.DoFnInfo; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.checkerframework.checker.nullness.qual.Nullable; /** * A {@link ParDoFnFactory} that creates a system {@link ParDoFn} responsible for transforming @@ -157,5 +159,10 @@ public void finishBundle() {} @Override public void abort() {} + + @Override + public @Nullable DoFnInfo getDoFnInfo() { + return null; + } } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ValuesDoFnFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ValuesDoFnFactory.java index acf143a5467b..ecf1d731a80d 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ValuesDoFnFactory.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ValuesDoFnFactory.java @@ -27,9 +27,11 @@ import org.apache.beam.runners.dataflow.worker.util.common.worker.Receiver; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.transforms.Values; +import org.apache.beam.sdk.util.DoFnInfo; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.WindowedValue; +import org.checkerframework.checker.nullness.qual.Nullable; /** * A {@link ParDoFnFactory} which returns a {@link ParDoFn} with similar behavior to {@link @@ -88,5 +90,10 @@ public void finishBundle() {} @Override public void abort() {} + + @Override + public @Nullable DoFnInfo getDoFnInfo() { + return null; + } } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoFn.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoFn.java index 86473ed0a337..d06de2be088e 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoFn.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoFn.java @@ -39,7 +39,9 @@ public interface ParDoFn { void abort() throws Exception; - default @Nullable DoFnInfo getDoFnInfo() throws Exception { - return null; - } + /** + * For ParDoFn's that wrap an underlying DoFn, returns a potentially cached copy of its DoFnInfo. + * Otherwise, returns null. + */ + @Nullable DoFnInfo getDoFnInfo() throws Exception; } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/SimplePartialGroupByKeyParDoFn.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/SimplePartialGroupByKeyParDoFn.java index cce77f3aceb6..eaf1f835a662 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/SimplePartialGroupByKeyParDoFn.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/SimplePartialGroupByKeyParDoFn.java @@ -17,6 +17,7 @@ */ package org.apache.beam.runners.dataflow.worker.util.common.worker; +import org.apache.beam.sdk.util.DoFnInfo; import org.checkerframework.checker.nullness.qual.Nullable; /** A partial group-by-key {@link ParDoFn} implementation. */ @@ -51,4 +52,9 @@ public void finishBundle() throws Exception { @Override public void abort() {} + + @Override + public @Nullable DoFnInfo getDoFnInfo() { + return null; + } } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/IntrinsicMapTaskExecutorTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/IntrinsicMapTaskExecutorTest.java index 73d69bd0f617..4d0d24a5886f 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/IntrinsicMapTaskExecutorTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/IntrinsicMapTaskExecutorTest.java @@ -69,7 +69,9 @@ import org.apache.beam.sdk.metrics.Metrics; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.util.DoFnInfo; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; +import org.checkerframework.checker.nullness.qual.Nullable; import org.junit.Assert; import org.junit.Rule; import org.junit.Test; @@ -223,6 +225,11 @@ public void abort() {} @Override public void finishKey(Object key) throws Exception {} + + @Override + public @Nullable DoFnInfo getDoFnInfo() { + return null; + } } /** Verify counts for the per-element-output-time counter are correct. */ diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/UserParDoFnFactoryTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/UserParDoFnFactoryTest.java index 43331d11a7ee..0168c8aa0ada 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/UserParDoFnFactoryTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/UserParDoFnFactoryTest.java @@ -196,7 +196,7 @@ public void testFactoryReuseInStep() throws Exception { parDoFn.startBundle(rcvr); parDoFn.processElement(WindowedValues.valueInGlobalWindow("foo")); - TestDoFn fn = (TestDoFn) ((SimpleParDoFn) parDoFn).getDoFnInfo().getDoFn(); + TestDoFn fn = (TestDoFn) ((SimpleParDoFn) parDoFn).getActiveDoFnInfo().getDoFn(); assertThat(fn, not(theInstance(initialFn))); parDoFn.finishBundle(); @@ -217,7 +217,8 @@ public void testFactoryReuseInStep() throws Exception { secondParDoFn.startBundle(rcvr); secondParDoFn.processElement(WindowedValues.valueInGlobalWindow("spam")); - TestDoFn reobtainedFn = (TestDoFn) ((SimpleParDoFn) secondParDoFn).getDoFnInfo().getDoFn(); + TestDoFn reobtainedFn = + (TestDoFn) ((SimpleParDoFn) secondParDoFn).getActiveDoFnInfo().getDoFn(); secondParDoFn.finishBundle(); assertThat(reobtainedFn.state, equalTo(TestDoFn.State.FINISHED)); @@ -256,13 +257,13 @@ public void testFactorySimultaneousUse() throws Exception { parDoFn.processElement(WindowedValues.valueInGlobalWindow("foo")); // Must be after the first call to process element for reallyStartBundle to have been called - TestDoFn firstDoFn = (TestDoFn) ((SimpleParDoFn) parDoFn).getDoFnInfo().getDoFn(); + TestDoFn firstDoFn = (TestDoFn) ((SimpleParDoFn) parDoFn).getActiveDoFnInfo().getDoFn(); secondParDoFn.startBundle(rcvr); secondParDoFn.processElement(WindowedValues.valueInGlobalWindow("spam")); // Must be after the first call to process element for reallyStartBundle to have been called - TestDoFn secondDoFn = (TestDoFn) ((SimpleParDoFn) secondParDoFn).getDoFnInfo().getDoFn(); + TestDoFn secondDoFn = (TestDoFn) ((SimpleParDoFn) secondParDoFn).getActiveDoFnInfo().getDoFn(); parDoFn.finishBundle(); secondParDoFn.finishBundle(); @@ -292,7 +293,7 @@ public void testFactoryDoesNotReuseAfterAborted() throws Exception { parDoFn.startBundle(rcvr); parDoFn.processElement(WindowedValues.valueInGlobalWindow("foo")); - TestDoFn fn = (TestDoFn) ((SimpleParDoFn) parDoFn).getDoFnInfo().getDoFn(); + TestDoFn fn = (TestDoFn) ((SimpleParDoFn) parDoFn).getActiveDoFnInfo().getDoFn(); parDoFn.abort(); assertThat(fn.state, equalTo(TestDoFn.State.TORN_DOWN)); @@ -310,7 +311,7 @@ public void testFactoryDoesNotReuseAfterAborted() throws Exception { secondParDoFn.startBundle(rcvr); secondParDoFn.processElement(WindowedValues.valueInGlobalWindow("foo")); - TestDoFn secondFn = (TestDoFn) ((SimpleParDoFn) secondParDoFn).getDoFnInfo().getDoFn(); + TestDoFn secondFn = (TestDoFn) ((SimpleParDoFn) secondParDoFn).getActiveDoFnInfo().getDoFn(); assertThat(secondFn, not(theInstance(fn))); assertThat(fn.state, equalTo(TestDoFn.State.TORN_DOWN)); diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutorTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutorTest.java index f707c0dd194f..f3c1bcacfc89 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutorTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutorTest.java @@ -67,7 +67,9 @@ import org.apache.beam.sdk.metrics.Metrics; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.util.DoFnInfo; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; +import org.checkerframework.checker.nullness.qual.Nullable; import org.junit.Assert; import org.junit.Rule; import org.junit.Test; @@ -221,6 +223,11 @@ public void abort() {} @Override public void finishKey(Object key) throws Exception {} + + @Override + public @Nullable DoFnInfo getDoFnInfo() { + return null; + } } /** Verify counts for the per-element-output-time counter are correct. */ diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoOperationTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoOperationTest.java index 0c9667acf38b..10203834d24f 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoOperationTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoOperationTest.java @@ -32,6 +32,8 @@ import org.apache.beam.runners.dataflow.worker.counters.CounterFactory.LongCounterMean; import org.apache.beam.runners.dataflow.worker.counters.CounterSet; import org.apache.beam.runners.dataflow.worker.counters.NameContext; +import org.apache.beam.sdk.util.DoFnInfo; +import org.checkerframework.checker.nullness.qual.Nullable; import org.hamcrest.CoreMatchers; import org.junit.Test; import org.junit.runner.RunWith; @@ -88,6 +90,11 @@ public void abort() throws Exception { @Override public void finishKey(Object key) throws Exception {} + + @Override + public @Nullable DoFnInfo getDoFnInfo() { + return null; + } } @Test From 11e08bbb9948a7f15eebbed6fc701d7d7c784975 Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Wed, 28 Oct 2026 14:14:18 -0700 Subject: [PATCH 9/9] Handle multi-output SDF's. --- .../IntrinsicMapTaskExecutorFactory.java | 13 +- ...treamingKeyedWorkItemSideInputParDoFn.java | 7 +- .../worker/StreamingModeExecutionContext.java | 53 ++++--- .../util/common/worker/ParDoOperation.java | 25 +-- .../ComputationWorkExecutorFactory.java | 76 ++++----- .../IntrinsicMapTaskExecutorFactoryTest.java | 105 ++++++++++++- .../StreamingModeExecutionContextTest.java | 94 ++++++++++- .../common/worker/MapTaskExecutorTest.java | 19 ++- .../ComputationWorkExecutorFactoryTest.java | 148 +++++++++++++++++- 9 files changed, 447 insertions(+), 93 deletions(-) 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 1009cf8e0406..ee3e6ac691a3 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 @@ -287,16 +287,19 @@ private OperationNode createParDoOperation( // An SDF's SplittableProcessFn behaves like a source, so we need to report its bytes read even // though it's not a read operator for autoscaling to work properly. - Coder bytesReadCoder = null; + ImmutableMap.Builder> bytesReadCodersBuilder = ImmutableMap.builder(); if (isSplittable(fn)) { - Node firstOutput = Iterables.getFirst(network.successors(node), null); - if (firstOutput instanceof OutputReceiverNode) { - bytesReadCoder = ((OutputReceiverNode) firstOutput).getCoder(); + for (Node successor : network.successors(node)) { + if (successor instanceof OutputReceiverNode) { + OutputReceiverNode outputReceiverNode = (OutputReceiverNode) successor; + bytesReadCodersBuilder.put( + outputReceiverNode.getOutputReceiver(), outputReceiverNode.getCoder()); + } } } OutputReceiver[] receivers = getOutputReceivers(network, node); return OperationNode.create( - new ParDoOperation(fn, receivers, operationContext, bytesReadCoder)); + new ParDoOperation(fn, receivers, operationContext, bytesReadCodersBuilder.build())); } private static boolean isSplittable(ParDoFn fn) throws Exception { diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputParDoFn.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputParDoFn.java index 3f9b695f89a5..c0081e35e4bf 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputParDoFn.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputParDoFn.java @@ -244,10 +244,11 @@ protected void onProcessWindowedValue(WindowedValue> el } /** - * Returns the {@link DoFnInfo} currently being used by this {@link SimpleParDoFn}. + * Returns the {@link DoFnInfo} currently being used by this {@link + * StreamingKeyedWorkItemSideInputParDoFn}. * - *

May be null if no element has been processed yet, or if the {@link SimpleParDoFn} has - * finished. + *

May be null if no element has been processed yet, or if the {@link + * StreamingKeyedWorkItemSideInputParDoFn} has finished. */ @VisibleForTesting @Nullable DoFnInfo getActiveDoFnInfo() { diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java index 3c37aae7306c..b400ad61e44e 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java @@ -48,6 +48,7 @@ import org.apache.beam.runners.core.metrics.ExecutionStateTracker; import org.apache.beam.runners.core.metrics.ExecutionStateTracker.ExecutionState; import org.apache.beam.runners.dataflow.worker.DataflowOperationContext.DataflowExecutionState; +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.NameContext; import org.apache.beam.runners.dataflow.worker.profiler.ScopedProfiler.ProfileScope; @@ -67,9 +68,10 @@ import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputStateFetcherFactory; import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor; import org.apache.beam.runners.dataflow.worker.util.common.worker.ElementCounter; +import org.apache.beam.runners.dataflow.worker.util.common.worker.Operation; import org.apache.beam.runners.dataflow.worker.util.common.worker.OutputObjectAndByteCounter; +import org.apache.beam.runners.dataflow.worker.util.common.worker.OutputReceiver; import org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoOperation; -import org.apache.beam.runners.dataflow.worker.util.common.worker.ReadOperation; import org.apache.beam.runners.dataflow.worker.util.common.worker.WorkExecutor; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.GlobalDataId; @@ -808,35 +810,38 @@ long computeSourceBytesProcessed(String sourceBytesCounterName) { // calls getAndReset on all modified counters to report to Dataflow. We should find a more // dedicated counter that is only extracted here. DataflowMapTaskExecutor mapTaskExecutor = (DataflowMapTaskExecutor) workExecutor; - long totalBytes = 0L; - ReadOperation readOperation = mapTaskExecutor.getReadOperation(); - if (readOperation != null - && readOperation.receivers != null - && readOperation.receivers.length > 0 - && readOperation.receivers[0] != null) { - HashMap readOpCounters = - readOperation.receivers[0].getOutputCounters(); - totalBytes += - Optional.ofNullable(readOpCounters.get(sourceBytesCounterName)) - .map(counter -> ((OutputObjectAndByteCounter) counter).getByteCount().getAndReset()) - .orElse(0L); - } + long totalBytes = + getAndResetOperationBytesProcessed( + mapTaskExecutor.getReadOperation(), sourceBytesCounterName); for (ParDoOperation parDoOp : mapTaskExecutor.getParDoOperations()) { - if (parDoOp.shouldReportBytesRead() - && parDoOp.receivers != null - && parDoOp.receivers.length > 0 - && parDoOp.receivers[0] != null) { - HashMap splittableProcessOpCounters = - parDoOp.receivers[0].getOutputCounters(); - totalBytes += - Optional.ofNullable(splittableProcessOpCounters.get(sourceBytesCounterName)) - .map(counter -> ((OutputObjectAndByteCounter) counter).getByteCount().getAndReset()) - .orElse(0L); + if (parDoOp.shouldReportBytesRead()) { + totalBytes += getAndResetOperationBytesProcessed(parDoOp, sourceBytesCounterName); } } return totalBytes; } + private static long getAndResetOperationBytesProcessed( + @Nullable Operation operation, String sourceBytesCounterName) { + if (operation == null || operation.receivers == null) { + return 0L; + } + long bytes = 0L; + for (OutputReceiver receiver : operation.receivers) { + if (receiver == null) { + continue; + } + ElementCounter counter = receiver.getOutputCounters().get(sourceBytesCounterName); + if (counter instanceof OutputObjectAndByteCounter) { + Counter byteCount = ((OutputObjectAndByteCounter) counter).getByteCount(); + if (byteCount != null) { + bytes += byteCount.getAndReset(); + } + } + } + return bytes; + } + public boolean advance() throws CoderException { if (!multiKeyBundleOptions.multiKeyBundleEnabled()) { return false; diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoOperation.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoOperation.java index 224a3b39577a..5376beb1829a 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoOperation.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoOperation.java @@ -18,29 +18,32 @@ package org.apache.beam.runners.dataflow.worker.util.common.worker; import java.io.Closeable; +import java.util.Map; import org.apache.beam.sdk.coders.Coder; 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.ImmutableMap; import org.checkerframework.checker.nullness.qual.Nullable; /** A ParDo mapping function. */ public class ParDoOperation extends ReceivingOperation { private final ParDoFn fn; - // If present, this ParDo should report the amount of bytes read using this coder to serialize - // elements. - private final @Nullable Coder bytesReadCoder; + // If non-empty, this ParDo should report the amount of bytes read using these per-output-receiver + // coders to serialize elements. + private final Map> bytesReadCoders; public ParDoOperation(ParDoFn fn, OutputReceiver[] outputReceivers, OperationContext context) { - this(fn, outputReceivers, context, null); + this(fn, outputReceivers, context, ImmutableMap.of()); } public ParDoOperation( ParDoFn fn, OutputReceiver[] outputReceivers, OperationContext context, - @Nullable Coder bytesReadCoder) { + @Nullable Map> bytesReadCoders) { super(outputReceivers, context); this.fn = fn; - this.bytesReadCoder = bytesReadCoder; + this.bytesReadCoders = + bytesReadCoders != null ? ImmutableMap.copyOf(bytesReadCoders) : ImmutableMap.of(); } @Override @@ -97,10 +100,14 @@ public ParDoFn getFn() throws Exception { } public boolean shouldReportBytesRead() { - return bytesReadCoder != null; + return !bytesReadCoders.isEmpty(); } - public @Nullable Coder getBytesReadCoder() { - return bytesReadCoder; + public Map> getBytesReadCoders() { + return bytesReadCoders; + } + + public @Nullable Coder getBytesReadCoder(OutputReceiver receiver) { + return bytesReadCoders.get(receiver); } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java index 8b3541f9b9c3..81ece67b74c8 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java @@ -53,7 +53,6 @@ import org.apache.beam.runners.dataflow.worker.streaming.harness.StreamingCounters; import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputStateFetcherFactory; import org.apache.beam.runners.dataflow.worker.util.common.worker.MapTaskExecutor; -import org.apache.beam.runners.dataflow.worker.util.common.worker.Operation; import org.apache.beam.runners.dataflow.worker.util.common.worker.OutputObjectAndByteCounter; import org.apache.beam.runners.dataflow.worker.util.common.worker.OutputReceiver; import org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoOperation; @@ -169,30 +168,23 @@ private static void trackAutoscalingBytesRead( MapTask mapTask, Nodes.ParallelInstructionNode readNode, Coder readCoder, - Operation operation, + OutputReceiver receiver, MapTaskExecutor mapTaskExecutor, String counterName) { - if (operation.receivers != null && operation.receivers.length > 0) { - NameContext nameContext = - NameContext.create( - mapTask.getStageName(), - readNode.getParallelInstruction().getOriginalName(), - readNode.getParallelInstruction().getSystemName(), - readNode.getParallelInstruction().getName()); - for (OutputReceiver receiver : operation.receivers) { - if (receiver != null) { - receiver.addOutputCounter( - counterName, - new OutputObjectAndByteCounter( - new IntrinsicMapTaskExecutorFactory.ElementByteSizeObservableCoder<>( - readCoder), - mapTaskExecutor.getOutputCounters(), - nameContext) - .setSamplingPeriod(100) - .countBytes(counterName)); - } - } - } + NameContext nameContext = + NameContext.create( + mapTask.getStageName(), + readNode.getParallelInstruction().getOriginalName(), + readNode.getParallelInstruction().getSystemName(), + readNode.getParallelInstruction().getName()); + receiver.addOutputCounter( + counterName, + new OutputObjectAndByteCounter( + new IntrinsicMapTaskExecutorFactory.ElementByteSizeObservableCoder<>(readCoder), + mapTaskExecutor.getOutputCounters(), + nameContext) + .setSamplingPeriod(100) + .countBytes(counterName)); } private static ReadOperation getValidatedReadOperation(MapTaskExecutor mapTaskExecutor) { @@ -237,28 +229,38 @@ ComputationWorkExecutor createComputationWorkExecutor( // If using a custom source, count bytes read for autoscaling. if (isCustomSource(readNode)) { - trackAutoscalingBytesRead( - mapTask, - readNode, - readCoder, - readOperation, - mapTaskExecutor, - computationState.sourceBytesProcessCounterName()); - } else { - for (ParDoOperation parDoOp : mapTaskExecutor.getParDoOperations()) { - if (parDoOp.shouldReportBytesRead()) { - Coder bytesReadCoder = parDoOp.getBytesReadCoder(); - if (bytesReadCoder != null) { + if (readOperation.receivers != null) { + for (OutputReceiver receiver : readOperation.receivers) { + if (receiver != null) { trackAutoscalingBytesRead( mapTask, readNode, - bytesReadCoder, - parDoOp, + readCoder, + receiver, mapTaskExecutor, computationState.sourceBytesProcessCounterName()); } } } + } else { + for (ParDoOperation parDoOp : mapTaskExecutor.getParDoOperations()) { + if (parDoOp.shouldReportBytesRead() && parDoOp.receivers != null) { + for (OutputReceiver receiver : parDoOp.receivers) { + if (receiver != null) { + Coder bytesReadCoder = parDoOp.getBytesReadCoder(receiver); + if (bytesReadCoder != null) { + trackAutoscalingBytesRead( + mapTask, + readNode, + bytesReadCoder, + receiver, + mapTaskExecutor, + computationState.sourceBytesProcessCounterName()); + } + } + } + } + } } ComputationWorkExecutor.Builder executionStateBuilder = 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 13c1a46704f0..83e2495a0839 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 @@ -110,6 +110,7 @@ import org.apache.beam.sdk.values.WindowedValues.FullWindowedValueCoder; import org.apache.beam.sdk.values.WindowingStrategy; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; 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.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.graph.MutableNetwork; @@ -573,7 +574,8 @@ public void testCreateParDoOperation() throws Exception { assertEquals(0, parDoOperation.receivers[0].getReceiverCount()); assertEquals(Operation.InitializationState.UNSTARTED, parDoOperation.initializationState); assertFalse(parDoOperation.shouldReportBytesRead()); - assertNull(parDoOperation.getBytesReadCoder()); + assertTrue(parDoOperation.getBytesReadCoders().isEmpty()); + assertNull(parDoOperation.getBytesReadCoder(parDoOperation.receivers[0])); } static class TestSplittableDoFn extends DoFn { @@ -658,7 +660,106 @@ public void testCreateSplittableProcessFnParDoOperation() throws Exception { assertTrue(parDoOperation.shouldReportBytesRead()); assertEquals( - WindowedValues.getValueOnlyCoder(StringUtf8Coder.of()), parDoOperation.getBytesReadCoder()); + ImmutableMap.of( + parDoOperation.receivers[0], WindowedValues.getValueOnlyCoder(StringUtf8Coder.of())), + parDoOperation.getBytesReadCoders()); + assertEquals( + WindowedValues.getValueOnlyCoder(StringUtf8Coder.of()), + parDoOperation.getBytesReadCoder(parDoOperation.receivers[0])); + } + + @Test + public void testCreateMultiOutputSplittableProcessFnParDoOperation() throws Exception { + int producerIndex = 1; + int producerOutputNum = 2; + + BatchModeExecutionContext context = + BatchModeExecutionContext.forTesting(options, counterSet, "testStage"); + TupleTag mainOutput = new TupleTag<>("main"); + TupleTag sideOutput = new TupleTag<>("side"); + String serializedFn = + StringUtils.byteArrayToJsonString( + SerializableUtils.serializeToByteArray( + DoFnInfo.forFn( + new TestSplittableDoFn(), + WindowingStrategy.globalDefault(), + null /* side input views */, + StringUtf8Coder.of() /* input coder */, + ImmutableMap.of( + mainOutput, StringUtf8Coder.of(), sideOutput, BigEndianIntegerCoder.of()), + mainOutput /* main output id */, + DoFnSchemaInformation.create(), + Collections.emptyMap()))); + + CloudObject cloudUserFn = CloudObject.forClassName("SplittableProcessFn"); + addString(cloudUserFn, PropertyNames.SERIALIZED_FN, serializedFn); + addObject( + cloudUserFn, + WorkerPropertyNames.RESTRICTION_CODER, + CloudObjects.asCloudObject( + KvCoder.of(OffsetRange.Coder.of(), VoidCoder.of()), /* sdkComponents= */ null)); + + MultiOutputInfo mainOutputInfo = new MultiOutputInfo(); + mainOutputInfo.setTag("main"); + MultiOutputInfo sideOutputInfo = new MultiOutputInfo(); + sideOutputInfo.setTag("side"); + + InstructionOutput mainCloudOutput = new InstructionOutput(); + mainCloudOutput.setName("SplittableParDo_main"); + mainCloudOutput.setCodec(windowedStringCoder); + mainCloudOutput.setOriginalName("originalName"); + mainCloudOutput.setSystemName("systemName"); + + Coder windowedIntCoder = WindowedValues.getValueOnlyCoder(BigEndianIntegerCoder.of()); + InstructionOutput sideCloudOutput = new InstructionOutput(); + sideCloudOutput.setName("SplittableParDo_side"); + sideCloudOutput.setCodec( + CloudObjects.asCloudObject(windowedIntCoder, /* sdkComponents= */ null)); + sideCloudOutput.setOriginalName("originalName"); + sideCloudOutput.setSystemName("systemName"); + + ParallelInstruction instruction = + createParDoInstruction( + producerIndex, producerOutputNum, "SplittableParDo", "", new TestSplittableDoFn()); + instruction.getParDo().setUserFn(cloudUserFn); + instruction.getParDo().setNumOutputs(2); + instruction.getParDo().setMultiOutputInfos(ImmutableList.of(mainOutputInfo, sideOutputInfo)); + instruction.setOutputs(ImmutableList.of(mainCloudOutput, sideCloudOutput)); + + ParallelInstructionNode instructionNode = + ParallelInstructionNode.create(instruction, ExecutionLocation.UNKNOWN); + + Function outputReceiversTransform = + IntrinsicMapTaskExecutorFactory.createOutputReceiversTransform(STAGE, counterSet, false); + Node mainOutputReceiverNode = + outputReceiversTransform.apply( + InstructionOutputNode.create(mainCloudOutput, PCOLLECTION_ID + "_main")); + Node sideOutputReceiverNode = + outputReceiversTransform.apply( + InstructionOutputNode.create(sideCloudOutput, PCOLLECTION_ID + "_side")); + + when(network.successors(instructionNode)) + .thenReturn(ImmutableSet.of(mainOutputReceiverNode, sideOutputReceiverNode)); + when(network.outDegree(instructionNode)).thenReturn(2); + when(network.edgesConnecting(instructionNode, mainOutputReceiverNode)) + .thenReturn(ImmutableSet.of(MultiOutputInfoEdge.create(mainOutputInfo))); + when(network.edgesConnecting(instructionNode, sideOutputReceiverNode)) + .thenReturn(ImmutableSet.of(MultiOutputInfoEdge.create(sideOutputInfo))); + + ArrayList createdOperations = new ArrayList<>(); + Node operationNode = + mapTaskExecutorFactory + .createOperationTransformForParallelInstructionNodes( + STAGE, network, options, readerRegistry, sinkRegistry, context, createdOperations) + .apply(instructionNode); + ParDoOperation parDoOperation = (ParDoOperation) ((OperationNode) operationNode).getOperation(); + + assertEquals(2, parDoOperation.receivers.length); + assertTrue(parDoOperation.shouldReportBytesRead()); + assertEquals( + WindowedValues.getValueOnlyCoder(StringUtf8Coder.of()), + parDoOperation.getBytesReadCoder(parDoOperation.receivers[0])); + assertEquals(windowedIntCoder, parDoOperation.getBytesReadCoder(parDoOperation.receivers[1])); } static ParallelInstruction createPartialGroupByKeyInstruction( diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java index 99c598e56088..bc808865cc68 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java @@ -102,6 +102,7 @@ import org.apache.beam.sdk.values.PCollectionView; import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; import org.hamcrest.Matchers; import org.joda.time.Duration; @@ -1007,7 +1008,10 @@ public void testComputeSourceBytesProcessed_splittableDoFn() throws Exception { ParDoOperation parDoOp = new ParDoOperation( - mockParDoFn, new OutputReceiver[] {parDoReceiver}, opContext, mock(Coder.class)); + mockParDoFn, + new OutputReceiver[] {parDoReceiver}, + opContext, + ImmutableMap.of(parDoReceiver, mock(Coder.class))); workExecutor = IntrinsicMapTaskExecutor.forOperations( @@ -1029,6 +1033,84 @@ public void testComputeSourceBytesProcessed_splittableDoFn() throws Exception { assertEquals(0L, executionContext.computeSourceBytesProcessed("sourceBytesProcessCounterName")); } + @Test + public void testComputeSourceBytesProcessed_multipleOutputReceivers() throws Exception { + NameContext nameContext = NameContext.create("stage", "original", "system", "out"); + DataflowOperationContext opContext = executionContext.createOperationContext(nameContext); + + // Custom source ReadOperation with a null receiver[0] and an active receiver[1]. + OutputReceiver readSideReceiver = new OutputReceiver(); + OutputObjectAndByteCounter readSideCounter = + new OutputObjectAndByteCounter( + mock(ElementByteSizeObservable.class), new CounterSet(), nameContext) + .countBytes("sourceBytesProcessCounterName"); + readSideCounter.getByteCount().addValue(15L); + readSideReceiver.addOutputCounter("sourceBytesProcessCounterName", readSideCounter); + ReadOperation readOp = + ReadOperation.create( + mock(NativeReader.class), new OutputReceiver[] {null, readSideReceiver}, opContext); + + // Multi-output SDF ParDoOperation with two receivers sharing a CounterSet (like production) + // plus an additional receiver with its own CounterSet. + ParDoFn mockParDoFn = mock(ParDoFn.class); + CounterSet sharedCounterSet = new CounterSet(); + + OutputReceiver sdfMainReceiver = new OutputReceiver(); + OutputObjectAndByteCounter sdfMainCounter = + new OutputObjectAndByteCounter( + mock(ElementByteSizeObservable.class), sharedCounterSet, nameContext) + .countBytes("sourceBytesProcessCounterName"); + sdfMainCounter.getByteCount().addValue(25L); + sdfMainReceiver.addOutputCounter("sourceBytesProcessCounterName", sdfMainCounter); + + OutputReceiver sdfSideReceiver1 = new OutputReceiver(); + OutputObjectAndByteCounter sdfSideCounter1 = + new OutputObjectAndByteCounter( + mock(ElementByteSizeObservable.class), sharedCounterSet, nameContext) + .countBytes("sourceBytesProcessCounterName"); + sdfSideCounter1.getByteCount().addValue(30L); + sdfSideReceiver1.addOutputCounter("sourceBytesProcessCounterName", sdfSideCounter1); + + OutputReceiver sdfSideReceiver2 = new OutputReceiver(); + OutputObjectAndByteCounter sdfSideCounter2 = + new OutputObjectAndByteCounter( + mock(ElementByteSizeObservable.class), new CounterSet(), nameContext) + .countBytes("sourceBytesProcessCounterName"); + sdfSideCounter2.getByteCount().addValue(35L); + sdfSideReceiver2.addOutputCounter("sourceBytesProcessCounterName", sdfSideCounter2); + + ParDoOperation sdfParDoOp = + new ParDoOperation( + mockParDoFn, + new OutputReceiver[] {sdfMainReceiver, null, sdfSideReceiver1, sdfSideReceiver2}, + opContext, + ImmutableMap.of( + sdfMainReceiver, + mock(Coder.class), + sdfSideReceiver1, + mock(Coder.class), + sdfSideReceiver2, + mock(Coder.class))); + + workExecutor = + IntrinsicMapTaskExecutor.forOperations( + Arrays.asList(readOp, sdfParDoOp), mock(ExecutionStateTracker.class)); + + Work work = + createMockWork( + Windmill.WorkItem.newBuilder() + .setKey(ByteString.copyFromUtf8("key1")) + .setWorkToken(1L) + .build(), + Watermarks.builder().setInputDataWatermark(Instant.EPOCH).build()); + + start(work); + + long processed = executionContext.computeSourceBytesProcessed("sourceBytesProcessCounterName"); + assertEquals(15L + 25L + 30L + 35L, processed); + assertEquals(0L, executionContext.computeSourceBytesProcessed("sourceBytesProcessCounterName")); + } + @Test public void testComputeSourceBytesProcessed_multipleParDoOperationsAndFlushState() throws Exception { @@ -1049,7 +1131,10 @@ public void testComputeSourceBytesProcessed_multipleParDoOperationsAndFlushState sdfReceiver1.addOutputCounter("sourceBytesProcessCounterName", sdfCounter1); ParDoOperation sdfParDoOp1 = new ParDoOperation( - mockParDoFn, new OutputReceiver[] {sdfReceiver1}, opContext, mock(Coder.class)); + mockParDoFn, + new OutputReceiver[] {sdfReceiver1}, + opContext, + ImmutableMap.of(sdfReceiver1, mock(Coder.class))); OutputReceiver sdfReceiver2 = new OutputReceiver(); OutputObjectAndByteCounter sdfCounter2 = @@ -1060,7 +1145,10 @@ public void testComputeSourceBytesProcessed_multipleParDoOperationsAndFlushState sdfReceiver2.addOutputCounter("sourceBytesProcessCounterName", sdfCounter2); ParDoOperation sdfParDoOp2 = new ParDoOperation( - mockParDoFn, new OutputReceiver[] {sdfReceiver2}, opContext, mock(Coder.class)); + mockParDoFn, + new OutputReceiver[] {sdfReceiver2}, + opContext, + ImmutableMap.of(sdfReceiver2, mock(Coder.class))); OutputReceiver regularReceiver = new OutputReceiver(); OutputObjectAndByteCounter regularCounter = diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutorTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutorTest.java index f3c1bcacfc89..48f0a7d1bfae 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutorTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutorTest.java @@ -32,6 +32,7 @@ import static org.hamcrest.Matchers.contains; import static org.hamcrest.Matchers.equalTo; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; import static org.mockito.ArgumentMatchers.anyBoolean; import static org.mockito.ArgumentMatchers.eq; @@ -68,6 +69,7 @@ import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.util.DoFnInfo; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; import org.checkerframework.checker.nullness.qual.Nullable; import org.junit.Assert; @@ -588,21 +590,23 @@ public void testGetParDoOperations() throws Exception { TestOutputReceiver receiver = new TestOutputReceiver(counterSet, NameContextsForTests.nameContextForTest()); ReadOperation readOp = new TestReadOperation(receiver, createContext("ReadOperation")); + OutputReceiver reportingReceiver1 = new OutputReceiver(); + OutputReceiver reportingReceiver2 = new OutputReceiver(); ParDoOperation regularParDo = new ParDoOperation( new NoopParDoFn(), new OutputReceiver[0], createContext("RegularParDo"), null); ParDoOperation reportingParDo1 = new ParDoOperation( new NoopParDoFn(), - new OutputReceiver[0], + new OutputReceiver[] {reportingReceiver1}, createContext("ReportingParDo1"), - StringUtf8Coder.of()); + ImmutableMap.of(reportingReceiver1, StringUtf8Coder.of())); ParDoOperation reportingParDo2 = new ParDoOperation( new NoopParDoFn(), - new OutputReceiver[0], + new OutputReceiver[] {reportingReceiver2}, createContext("ReportingParDo2"), - VarIntCoder.of()); + ImmutableMap.of(reportingReceiver2, VarIntCoder.of())); List operations = Arrays.asList(readOp, regularParDo, reportingParDo1, reportingParDo2); @@ -610,9 +614,10 @@ public void testGetParDoOperations() throws Exception { try (MapTaskExecutor executor = new MapTaskExecutor(operations, counterSet, stateTracker)) { List parDoOps = executor.getParDoOperations(); assertThat(parDoOps, contains(regularParDo, reportingParDo1, reportingParDo2)); - Assert.assertNull(parDoOps.get(0).getBytesReadCoder()); - assertEquals(StringUtf8Coder.of(), parDoOps.get(1).getBytesReadCoder()); - assertEquals(VarIntCoder.of(), parDoOps.get(2).getBytesReadCoder()); + assertTrue(parDoOps.get(0).getBytesReadCoders().isEmpty()); + Assert.assertNull(parDoOps.get(0).getBytesReadCoder(reportingReceiver1)); + assertEquals(StringUtf8Coder.of(), parDoOps.get(1).getBytesReadCoder(reportingReceiver1)); + assertEquals(VarIntCoder.of(), parDoOps.get(2).getBytesReadCoder(reportingReceiver2)); } } } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactoryTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactoryTest.java index 5c0adf1e7cbd..4022ea2c62da 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactoryTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactoryTest.java @@ -22,6 +22,7 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; import static org.mockito.Mockito.mock; import com.google.api.services.dataflow.model.InstructionInput; @@ -80,6 +81,7 @@ import org.apache.beam.sdk.values.WindowedValues; import org.apache.beam.sdk.values.WindowingStrategy; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -288,13 +290,13 @@ private static ParallelInstruction createParDoInstruction( ParDoOperation stringSplittableOp = allParDoOps.stream() - .filter(op -> windowedStringCoder.equals(op.getBytesReadCoder())) + .filter(op -> windowedStringCoder.equals(op.getBytesReadCoder(op.receivers[0]))) .findFirst() .orElseThrow( () -> new AssertionError("Expected ParDoOperation with windowedStringCoder")); ParDoOperation intSplittableOp = allParDoOps.stream() - .filter(op -> windowedIntCoder.equals(op.getBytesReadCoder())) + .filter(op -> windowedIntCoder.equals(op.getBytesReadCoder(op.receivers[0]))) .findFirst() .orElseThrow(() -> new AssertionError("Expected ParDoOperation with windowedIntCoder")); @@ -308,7 +310,8 @@ private static ParallelInstruction createParDoInstruction( .filter(op -> !op.shouldReportBytesRead()) .findFirst() .orElseThrow(() -> new AssertionError("Expected non-reporting ParDoOperation")); - assertNull(regularParDoOp.getBytesReadCoder()); + assertTrue(regularParDoOp.getBytesReadCoders().isEmpty()); + assertNull(regularParDoOp.getBytesReadCoder(regularParDoOp.receivers[0])); assertNull(regularParDoOp.receivers[0].getOutputCounters().get(counterName)); // Process elements through each operation's receiver and verify each uses its own coder. @@ -330,4 +333,143 @@ private static ParallelInstruction createParDoInstruction( assertEquals( Long.valueOf(expectedStringBytes + expectedIntBytes), bytesReadCounter.getAggregate()); } + + @Test + public void testCreateComputationWorkExecutor_multiOutputSplittableParDoWithDifferentCoders() + throws Exception { + DataflowWorkerHarnessOptions options = + PipelineOptionsFactory.as(DataflowWorkerHarnessOptions.class); + ComputationWorkExecutorFactory factory = + new ComputationWorkExecutorFactory( + options, + IntrinsicMapTaskExecutorFactory.defaultFactory(), + mock(ReaderCache.class), + computationId -> mock(WindmillStateCache.ForComputation.class), + DataflowExecutionStateSampler.instance(), + StreamingCounters.create(), + mock(FailureTracker.class), + IdGenerators.decrementingLongs(), + new FakeGlobalConfigHandle(StreamingGlobalConfig.builder().build()), + mock(HotKeyLogger.class), + SideInputStateFetcherFactory.fromOptions(options), + MultiKeyBundleOptions.fromOptions(options)); + + Coder> windowedStringCoder = + WindowedValues.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE); + Coder> windowedIntCoder = + WindowedValues.getFullCoder(VarIntCoder.of(), GlobalWindow.Coder.INSTANCE); + + TupleTag mainOutputTag = new TupleTag<>("main"); + TupleTag sideOutputTag = new TupleTag<>("side"); + String serializedFn = + StringUtils.byteArrayToJsonString( + SerializableUtils.serializeToByteArray( + DoFnInfo.forFn( + new StringSplittableDoFn(), + WindowingStrategy.globalDefault(), + null /* side input views */, + StringUtf8Coder.of() /* input coder */, + ImmutableMap.of( + mainOutputTag, StringUtf8Coder.of(), sideOutputTag, VarIntCoder.of()), + mainOutputTag, + DoFnSchemaInformation.create(), + Collections.emptyMap()))); + + CloudObject cloudUserFn = CloudObject.forClassName("SplittableProcessFn"); + addString(cloudUserFn, PropertyNames.SERIALIZED_FN, serializedFn); + addObject( + cloudUserFn, + WorkerPropertyNames.RESTRICTION_CODER, + CloudObjects.asCloudObject( + KvCoder.of(OffsetRange.Coder.of(), VoidCoder.of()), /* sdkComponents= */ null)); + + InstructionInput cloudInput = new InstructionInput(); + cloudInput.setProducerInstructionIndex(0); + cloudInput.setOutputNum(0); + + MultiOutputInfo mainOutputInfo = new MultiOutputInfo(); + mainOutputInfo.setTag("main"); + MultiOutputInfo sideOutputInfo = new MultiOutputInfo(); + sideOutputInfo.setTag("side"); + + ParDoInstruction parDoInstruction = new ParDoInstruction(); + parDoInstruction.setInput(cloudInput); + parDoInstruction.setNumOutputs(2); + parDoInstruction.setMultiOutputInfos(ImmutableList.of(mainOutputInfo, sideOutputInfo)); + parDoInstruction.setUserFn(cloudUserFn); + + InstructionOutput mainOutput = new InstructionOutput(); + mainOutput.setName("MultiOutputSplittableParDo_main"); + mainOutput.setCodec(CloudObjects.asCloudObject(windowedStringCoder, /* sdkComponents= */ null)); + mainOutput.setOriginalName("MultiOutputSplittableParDoOriginal"); + mainOutput.setSystemName("MultiOutputSplittableParDo"); + + InstructionOutput sideOutput = new InstructionOutput(); + sideOutput.setName("MultiOutputSplittableParDo_side"); + sideOutput.setCodec(CloudObjects.asCloudObject(windowedIntCoder, /* sdkComponents= */ null)); + sideOutput.setOriginalName("MultiOutputSplittableParDoOriginal"); + sideOutput.setSystemName("MultiOutputSplittableParDo"); + + ParallelInstruction multiOutputSdfInstruction = new ParallelInstruction(); + multiOutputSdfInstruction.setParDo(parDoInstruction); + multiOutputSdfInstruction.setOutputs(ImmutableList.of(mainOutput, sideOutput)); + multiOutputSdfInstruction.setName("MultiOutputSplittableParDo"); + multiOutputSdfInstruction.setSystemName("MultiOutputSplittableParDo"); + multiOutputSdfInstruction.setOriginalName("MultiOutputSplittableParDoOriginal"); + + MapTask mapTask = new MapTask(); + mapTask.setStageName("testStage"); + mapTask.setSystemName("testSystem"); + mapTask.setInstructions( + Arrays.asList( + createWindmillReadInstruction("Read", windowedStringCoder), multiOutputSdfInstruction)); + mapTask.setFactory(Transport.getJsonFactory()); + + StageInfo stageInfo = StageInfo.create("testStage", "testSystem"); + ComputationState computationState = + new ComputationState( + "computationId", + mapTask, + mock(BoundedQueueExecutor.class), + Collections.emptyMap(), + mock(WindmillStateCache.ForComputation.class)); + + ComputationWorkExecutor workExecutor = + factory.createComputationWorkExecutor(stageInfo, computationState, "workLatencyTrackingId"); + + DataflowMapTaskExecutor mapTaskExecutor = (DataflowMapTaskExecutor) workExecutor.workExecutor(); + List allParDoOps = mapTaskExecutor.getParDoOperations(); + assertEquals(1, allParDoOps.size()); + + ParDoOperation multiOutputSdfOp = allParDoOps.get(0); + assertTrue(multiOutputSdfOp.shouldReportBytesRead()); + assertEquals(2, multiOutputSdfOp.receivers.length); + assertEquals( + windowedStringCoder, multiOutputSdfOp.getBytesReadCoder(multiOutputSdfOp.receivers[0])); + assertEquals( + windowedIntCoder, multiOutputSdfOp.getBytesReadCoder(multiOutputSdfOp.receivers[1])); + + String counterName = computationState.sourceBytesProcessCounterName(); + assertNotNull(multiOutputSdfOp.receivers[0].getOutputCounters().get(counterName)); + assertNotNull(multiOutputSdfOp.receivers[1].getOutputCounters().get(counterName)); + + // Process differently typed elements on main and side output receivers without + // ClassCastException. + WindowedValue stringElem = WindowedValues.valueInGlobalWindow("hello"); + WindowedValue intElem = WindowedValues.valueInGlobalWindow(12345); + multiOutputSdfOp.receivers[0].process(stringElem); + multiOutputSdfOp.receivers[1].process(intElem); + + long expectedStringBytes = + CoderUtils.encodeToByteArray(windowedStringCoder, stringElem, Coder.Context.NESTED).length; + long expectedIntBytes = + CoderUtils.encodeToByteArray(windowedIntCoder, intElem, Coder.Context.NESTED).length; + + @SuppressWarnings("unchecked") + Counter bytesReadCounter = + (Counter) mapTaskExecutor.getOutputCounters().getExistingCounter(counterName); + assertNotNull(bytesReadCounter); + assertEquals( + Long.valueOf(expectedStringBytes + expectedIntBytes), bytesReadCounter.getAggregate()); + } }