From e7472da72152d7af5feb39d8198017a60adddab1 Mon Sep 17 00:00:00 2001 From: Reuven Lax Date: Wed, 29 Apr 2026 10:58:09 -0700 Subject: [PATCH 01/11] side-input improvements --- .../beam/runners/core/SimpleDoFnRunner.java | 26 +- .../worker/SimpleDoFnRunnerFactory.java | 12 - .../dataflow/worker/SimpleParDoFn.java | 78 ++- .../worker/StreamingSideInputDoFnRunner.java | 38 +- .../worker/StreamingSideInputFetcher.java | 32 +- .../worker/StreamingSideInputProcessor.java | 135 ++++ .../StreamingSideInputProcessorTest.java | 205 +++++++ runners/prism/java/build.gradle | 3 + .../beam/runners/samza/runtime/DoFnOp.java | 579 ++++++++++++++++++ .../runners/samza/runtime/PortableDoFnOp.java | 467 ++++++++++++++ .../sdk/testing/UsesSideInputsInTimer.java | 27 + .../org/apache/beam/sdk/transforms/DoFn.java | 16 + .../transforms/reflect/DoFnSignatures.java | 9 +- .../apache/beam/sdk/transforms/ParDoTest.java | 149 +++++ .../beam/fn/harness/FnApiDoFnRunner.java | 27 + 15 files changed, 1734 insertions(+), 69 deletions(-) create mode 100644 runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputProcessor.java create mode 100644 runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputProcessorTest.java create mode 100644 runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnOp.java create mode 100644 runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/PortableDoFnOp.java create mode 100644 sdks/java/core/src/main/java/org/apache/beam/sdk/testing/UsesSideInputsInTimer.java diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java index 9859d672b3e8..470e22a66991 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java @@ -870,8 +870,17 @@ public InputT element(DoFn doFn) { } @Override - public Object sideInput(String tagId) { - throw new UnsupportedOperationException("SideInput parameters are not supported."); + public @Nullable Object sideInput(String tagId) { + PCollectionView view = + checkStateNotNull(sideInputMapping.get(tagId), "Side input tag %s not found", tagId); + return sideInput(view); + } + + @Override + public T sideInput(PCollectionView view) { + checkNotNull(view, "View passed to sideInput cannot be null"); + return SimpleDoFnRunner.this.sideInput( + view, view.getWindowMappingFn().getSideInputWindow(window())); } @Override @@ -1196,8 +1205,17 @@ public InputT element(DoFn doFn) { } @Override - public Object sideInput(String tagId) { - throw new UnsupportedOperationException("SideInput parameters are not supported."); + public @Nullable Object sideInput(String tagId) { + PCollectionView view = + checkStateNotNull(sideInputMapping.get(tagId), "Side input tag %s not found", tagId); + return sideInput(view); + } + + @Override + public T sideInput(PCollectionView view) { + checkNotNull(view, "View passed to sideInput cannot be null"); + return SimpleDoFnRunner.this.sideInput( + view, view.getWindowMappingFn().getSideInputWindow(window())); } @Override diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleDoFnRunnerFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleDoFnRunnerFactory.java index 52fcec439aaf..5286fc1aae90 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleDoFnRunnerFactory.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleDoFnRunnerFactory.java @@ -24,7 +24,6 @@ import org.apache.beam.runners.core.SideInputReader; import org.apache.beam.sdk.coders.Coder; 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.DoFnSchemaInformation; import org.apache.beam.sdk.util.WindowedValueMultiReceiver; @@ -68,17 +67,6 @@ public DoFnRunner createRunner( windowingStrategy, doFnSchemaInformation, sideInputMapping); - boolean hasStreamingSideInput = - options.as(StreamingOptions.class).isStreaming() && !sideInputReader.isEmpty(); - if (hasStreamingSideInput) { - return new StreamingSideInputDoFnRunner<>( - fnRunner, - new StreamingSideInputFetcher<>( - sideInputViews, - inputCoder, - windowingStrategy, - (StreamingModeExecutionContext.StreamingModeStepContext) userStepContext)); - } return fnRunner; } } 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 434d46c20a5b..7203cbcae305 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 @@ -22,7 +22,9 @@ import java.io.Closeable; import java.util.Collection; +import java.util.Collections; import java.util.HashMap; +import java.util.Iterator; import java.util.List; import java.util.Map; import org.apache.beam.runners.core.DoFnRunner; @@ -59,6 +61,7 @@ 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.ImmutableList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Duration; import org.joda.time.Instant; @@ -76,7 +79,7 @@ "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) "nullness" // TODO(https://github.com/apache/beam/issues/20497) }) -public class SimpleParDoFn implements ParDoFn { +public class SimpleParDoFn implements ParDoFn { // TODO: Remove once Distributions has shipped. @VisibleForTesting @@ -112,6 +115,8 @@ public class SimpleParDoFn implements ParDoFn { // GroupAlsoByWindowViaWindowSetDoFn private @Nullable DoFnSignature fnSignature; + private @Nullable StreamingSideInputProcessor sideInputProcessor; + /** Creates a {@link SimpleParDoFn} using basic information about the step being executed. */ SimpleParDoFn( PipelineOptions options, @@ -317,8 +322,32 @@ public void output(TupleTag tag, WindowedValue output) { outputManager, doFnSchemaInformation, sideInputMapping); + if (hasStreamingSideInput) { + sideInputProcessor = + new StreamingSideInputProcessor<>( + new StreamingSideInputFetcher( + fnInfo.getSideInputViews(), + fnInfo.getInputCoder(), + (WindowingStrategy) fnInfo.getWindowingStrategy(), + (StreamingModeExecutionContext.StreamingModeStepContext) userStepContext)); + } fnRunner.startBundle(); + if (sideInputProcessor != null) { + boolean hasState = fnSignature != null && !fnSignature.stateDeclarations().isEmpty(); + Iterator> unblockedElements = sideInputProcessor.tryUnblockElements(); + for (Iterator> it = unblockedElements; it.hasNext(); ) { + WindowedValue unblockedElement = it.next(); + fnRunner.processElement(unblockedElement); + if (hasState) { + // These elements are now processed. Register cleanup timers for all the unblocked + // windows. + registerStateCleanup( + (WindowingStrategy) getDoFnInfo().getWindowingStrategy(), + (Collection) unblockedElement.getWindows()); + } + } + } } @Override @@ -334,14 +363,32 @@ public void processElement(Object untypedElem) throws Exception { WindowedValue elem = (WindowedValue) untypedElem; - if (fnSignature != null && fnSignature.stateDeclarations().size() > 0) { + boolean hasState = fnSignature != null && !fnSignature.stateDeclarations().isEmpty(); + outputsPerElementTracker.onProcessElement(); + + Collection windowsProcessed; + if (sideInputProcessor != null) { + windowsProcessed = hasState ? Lists.newArrayList() : Collections.emptyList(); + for (Iterator> it = + sideInputProcessor.handleProcessElement(elem); + it.hasNext(); ) { + WindowedValue toProcess = it.next(); + fnRunner.processElement(toProcess); + if (hasState) { + windowsProcessed.addAll((Collection) toProcess.getWindows()); + // If the element was blocked, don't register a cleanup timer. The timer will be + // registered + // when the window is unblocked ensuring that it is not processed until the element is. + } + } + } else { + fnRunner.processElement(elem); + windowsProcessed = (Collection) elem.getWindows(); + } + if (hasState) { registerStateCleanup( - (WindowingStrategy) getDoFnInfo().getWindowingStrategy(), - (Collection) elem.getWindows()); + (WindowingStrategy) getDoFnInfo().getWindowingStrategy(), windowsProcessed); } - - outputsPerElementTracker.onProcessElement(); - fnRunner.processElement(elem); outputsPerElementTracker.onProcessElementSuccess(); } @@ -367,6 +414,9 @@ private void processUserTimer(TimerData timer) throws Exception { if (fnSignature.timerDeclarations().containsKey(timer.getTimerId()) || fnSignature.timerFamilyDeclarations().containsKey(timer.getTimerFamilyId())) { BoundedWindow window = ((WindowNamespace) timer.getNamespace()).getWindow(); + if (sideInputProcessor != null) { + sideInputProcessor.handleProcessTimer(timer); + } fnRunner.onTimer( timer.getTimerId(), timer.getTimerFamilyId(), @@ -380,7 +430,6 @@ private void processUserTimer(TimerData timer) throws Exception { } private void processSystemTimer(TimerData timer) throws Exception { - // Timer owned by this class, for cleaning up state in expired windows if (timer.getTimerId().equals(CLEANUP_TIMER_ID)) { checkState( @@ -396,6 +445,13 @@ private void processSystemTimer(TimerData timer) throws Exception { WindowNamespace.class.getSimpleName(), timer); + if (sideInputProcessor != null) { + // We must call this to ensure the side-input is cached for onWindowExpiration. Since we + // don't set cleanup + // timers until we actually call processElement, the window must be unblocked here. + sideInputProcessor.handleProcessTimer(timer); + } + BoundedWindow window = ((WindowNamespace) timer.getNamespace()).getWindow(); Instant targetTime = earliestAllowableCleanupTime(window, fnInfo.getWindowingStrategy()); @@ -436,10 +492,14 @@ private void processSystemTimer(TimerData timer) throws Exception { public void finishBundle() throws Exception { if (fnRunner != null) { fnRunner.finishBundle(); + if (sideInputProcessor != null) { + sideInputProcessor.handleFinishBundle(); + } doFnInstanceManager.complete(fnInfo); fnRunner = null; fnInfo = null; fnSignature = null; + sideInputProcessor = null; } } @@ -490,7 +550,7 @@ private void processTimers( } } - private void registerStateCleanup( + private void registerStateCleanup( WindowingStrategy windowingStrategy, Collection windowsToCleanup) { Coder windowCoder = windowingStrategy.getWindowFn().windowCoder(); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputDoFnRunner.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputDoFnRunner.java index 3b7891c5378d..b41b0c5049a7 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputDoFnRunner.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputDoFnRunner.java @@ -17,9 +17,8 @@ */ package org.apache.beam.runners.dataflow.worker; -import java.util.Set; +import java.util.Iterator; import org.apache.beam.runners.core.DoFnRunner; -import org.apache.beam.sdk.state.BagState; import org.apache.beam.sdk.state.TimeDomain; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; @@ -37,43 +36,32 @@ public class StreamingSideInputDoFnRunner implements DoFnRunner { private final DoFnRunner simpleDoFnRunner; - private final StreamingSideInputFetcher sideInputFetcher; + private final StreamingSideInputProcessor sideInputProcessor; public StreamingSideInputDoFnRunner( DoFnRunner simpleDoFnRunner, StreamingSideInputFetcher sideInputFetcher) { this.simpleDoFnRunner = simpleDoFnRunner; - this.sideInputFetcher = sideInputFetcher; + this.sideInputProcessor = new StreamingSideInputProcessor<>(sideInputFetcher); } @Override public void startBundle() { simpleDoFnRunner.startBundle(); - sideInputFetcher.prefetchBlockedMap(); - - // Find the set of ready windows. - Set readyWindows = sideInputFetcher.getReadyWindows(); - - Iterable>> elementsBags = - sideInputFetcher.prefetchElements(readyWindows); - - // Run the DoFn code now that all side inputs are ready. - for (BagState> elementsBag : elementsBags) { - Iterable> elements = elementsBag.read(); - for (WindowedValue elem : elements) { - simpleDoFnRunner.processElement(elem); - } - elementsBag.clear(); + Iterator> unblocked = sideInputProcessor.tryUnblockElements(); + for (Iterator> it = unblocked; it.hasNext(); ) { + WindowedValue elem = it.next(); + simpleDoFnRunner.processElement(elem); } - sideInputFetcher.releaseBlockedWindows(readyWindows); } @Override public void processElement(WindowedValue compressedElem) { - for (WindowedValue elem : compressedElem.explodeWindows()) { - if (!sideInputFetcher.storeIfBlocked(elem)) { - simpleDoFnRunner.processElement(elem); - } + for (Iterator> it = + sideInputProcessor.handleProcessElement(compressedElem); + it.hasNext(); ) { + WindowedValue elem = it.next(); + simpleDoFnRunner.processElement(elem); } } @@ -94,7 +82,7 @@ public void onTimer( @Override public void finishBundle() { simpleDoFnRunner.finishBundle(); - sideInputFetcher.persist(); + sideInputProcessor.handleFinishBundle(); } @Override diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputFetcher.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputFetcher.java index 76913baa6aa7..e97e16ca3133 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputFetcher.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputFetcher.java @@ -20,6 +20,7 @@ import java.io.IOException; import java.io.InputStream; import java.io.OutputStream; +import java.util.Collections; import java.util.HashMap; import java.util.HashSet; import java.util.List; @@ -83,7 +84,6 @@ public StreamingSideInputFetcher( this.stepContext = stepContext; this.mainWindowCoder = windowingStrategy.getWindowFn().windowCoder(); - this.sideInputViews = new HashMap<>(); for (PCollectionView view : views) { sideInputViews.put(view.getTagInternal().getId(), view); @@ -188,11 +188,7 @@ public Iterable> prefetchTimers(Iterable readyWindows) { return timers; } - /** Compute the set of side inputs that are not yet ready for the given main input window. */ - public boolean storeIfBlocked(WindowedValue elem) { - @SuppressWarnings("unchecked") - W window = (W) Iterables.getOnlyElement(elem.getWindows()); - + private Set checkIfBlocked(W window) { Set blocked = blockedMap().get(window); if (blocked == null) { for (PCollectionView view : sideInputViews.values()) { @@ -205,7 +201,16 @@ public boolean storeIfBlocked(WindowedValue elem) { } } } - if (blocked != null) { + return blocked == null ? Collections.emptySet() : blocked; + } + + /** Compute the set of side inputs that are not yet ready for the given main input window. */ + public boolean storeIfBlocked(WindowedValue elem) { + @SuppressWarnings("unchecked") + W window = (W) Iterables.getOnlyElement(elem.getWindows()); + + Set blocked = checkIfBlocked(window); + if (!blocked.isEmpty()) { elementBag(window).add(elem); watermarkHold(window).add(elem.getTimestamp()); stepContext.addBlockingSideInputs(blocked); @@ -223,17 +228,12 @@ public boolean storeIfBlocked(TimerData timer) { @SuppressWarnings("unchecked") WindowNamespace windowNamespace = (WindowNamespace) timer.getNamespace(); W window = windowNamespace.getWindow(); - - boolean blocked = false; - for (PCollectionView view : sideInputViews.values()) { - if (!stepContext.issueSideInputFetch(view, window, SideInputState.UNKNOWN)) { - blocked = true; - } - } - if (blocked) { + Set blocked = checkIfBlocked(window); + if (!blocked.isEmpty()) { timerBag(window).add(timer); + return true; } - return blocked; + return false; } public void persist() { diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputProcessor.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputProcessor.java new file mode 100644 index 000000000000..f51312a6e9dc --- /dev/null +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputProcessor.java @@ -0,0 +1,135 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.dataflow.worker; + +import java.util.Iterator; +import java.util.NoSuchElementException; +import java.util.Set; +import org.apache.beam.runners.core.TimerInternals; +import org.apache.beam.sdk.state.BagState; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterators; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** Helper class for handling elements blocked on side inputs. */ +class StreamingSideInputProcessor { + private final StreamingSideInputFetcher sideInputFetcher; + + public StreamingSideInputProcessor(StreamingSideInputFetcher sideInputFetcher) { + this.sideInputFetcher = sideInputFetcher; + } + + /** + * Handle's startBundle. If there are unblocked elements, process them and then return the set of + * windows that were unblocked. + */ + Iterator> tryUnblockElements() { + sideInputFetcher.prefetchBlockedMap(); + + // Find the set of ready windows. + Set readyWindows = sideInputFetcher.getReadyWindows(); + + Iterable>> elementsBags = + sideInputFetcher.prefetchElements(readyWindows); + + // Return a lazy iterator to the released elements. This is a destructive iterator - it clears + // the bags after reading them. Bags can be paged in from the service, so we try to avoid + // materializing the whole + // bag into memory here. + Iterator> releasedElements = + new Iterator>() { + Iterator>> bagsIterator = elementsBags.iterator(); + @Nullable Iterator> currentBagElements; + @Nullable BagState> currentBag; + + @Override + public boolean hasNext() { + do { + if (currentBagElements == null || !currentBagElements.hasNext()) { + if (!advanceBag()) { + // We're done iterating - release the blocked windows. + sideInputFetcher.releaseBlockedWindows(readyWindows); + return false; + } + } + } while (!org.apache.beam.sdk.util.Preconditions.checkStateNotNull(currentBagElements) + .hasNext()); + return true; + } + + boolean advanceBag() { + // Once we finish reading a bag, clear it. + clearCurrentBag(); + if (bagsIterator.hasNext()) { + currentBag = bagsIterator.next(); + currentBagElements = currentBag.read().iterator(); + return true; + } else { + return false; + } + } + + void clearCurrentBag() { + if (currentBag != null) { + currentBag.clear(); + currentBag = null; + currentBagElements = null; + } + } + + @Override + public WindowedValue next() { + if (!hasNext()) { + throw new NoSuchElementException(); + } + return org.apache.beam.sdk.util.Preconditions.checkStateNotNull(currentBagElements) + .next(); + } + }; + + return releasedElements; + } + + void handleFinishBundle() { + sideInputFetcher.persist(); + } + + /* + Handle process element. Runs the elements that have an available side input, and buffers elements for which the + side input is blocked. Returns the list of elements that are unblocked and should be processed. + */ + Iterator> handleProcessElement( + WindowedValue compressedElem) { + // Note: We could write this as a three-line stream expression, but side effects are discouraged + // in Java streams. + return Iterators.filter( + compressedElem.explodeWindows().iterator(), + (WindowedValue e) -> !sideInputFetcher.storeIfBlocked(e)); + } + + void handleProcessTimer(TimerInternals.TimerData timer) { + // We must call this to ensure the side-input is cached for the timer. However since a user + // timer can only + // be set via element processing (or another timer) in the same window, the window should be + // unblocked once + // we get here. + Preconditions.checkState(!sideInputFetcher.storeIfBlocked(timer)); + } +} diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputProcessorTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputProcessorTest.java new file mode 100644 index 000000000000..804c22ea8567 --- /dev/null +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputProcessorTest.java @@ -0,0 +1,205 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.dataflow.worker; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.containsInAnyOrder; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anySet; +import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.google.api.client.util.Lists; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.Iterator; +import java.util.Set; +import org.apache.beam.runners.core.TimerInternals.TimerData; +import org.apache.beam.sdk.state.BagState; +import org.apache.beam.sdk.transforms.windowing.IntervalWindow; +import org.apache.beam.sdk.transforms.windowing.PaneInfo; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowedValues; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; +import org.joda.time.Instant; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; + +/** Unit tests for {@link StreamingSideInputProcessor}. */ +@RunWith(JUnit4.class) +public class StreamingSideInputProcessorTest { + + @Mock private StreamingSideInputFetcher mockFetcher; + private StreamingSideInputProcessor processor; + + @Before + public void setUp() { + MockitoAnnotations.initMocks(this); + processor = new StreamingSideInputProcessor<>(mockFetcher); + } + + @Test + public void testTryUnblockElementsNoReadyWindows() { + // Given + doNothing().when(mockFetcher).prefetchBlockedMap(); + when(mockFetcher.getReadyWindows()).thenReturn(Collections.emptySet()); + + // When + Iterator> unblocked = processor.tryUnblockElements(); + + // Then + assertFalse(unblocked.hasNext()); + verify(mockFetcher).prefetchBlockedMap(); + verify(mockFetcher).getReadyWindows(); + } + + @Test + public void testTryUnblockElementsWithReadyWindows() { + // Given + IntervalWindow window1 = new IntervalWindow(new Instant(0), new Instant(10)); + IntervalWindow window2 = new IntervalWindow(new Instant(10), new Instant(20)); + Set readyWindows = new HashSet<>(Arrays.asList(window1, window2)); + + WindowedValue element1 = + WindowedValues.of("e1", new Instant(5), Arrays.asList(window1), PaneInfo.NO_FIRING); + WindowedValue element2 = + WindowedValues.of("e2", new Instant(15), Arrays.asList(window2), PaneInfo.NO_FIRING); + + @SuppressWarnings("unchecked") + BagState> mockBag1 = mock(BagState.class); + @SuppressWarnings("unchecked") + BagState> mockBag2 = mock(BagState.class); + + when(mockBag1.read()).thenReturn(Arrays.asList(element1)); + when(mockBag2.read()).thenReturn(Arrays.asList(element2)); + + doNothing().when(mockFetcher).prefetchBlockedMap(); + when(mockFetcher.getReadyWindows()).thenReturn(readyWindows); + when(mockFetcher.prefetchElements(readyWindows)).thenReturn(Arrays.asList(mockBag1, mockBag2)); + doNothing().when(mockFetcher).releaseBlockedWindows(readyWindows); + + // When + Iterable> unblocked = () -> processor.tryUnblockElements(); + + // Then + verify(mockBag1, never()).clear(); + verify(mockFetcher, never()).releaseBlockedWindows(anySet()); + assertThat(unblocked, containsInAnyOrder(element1, element2)); + verify(mockFetcher).prefetchBlockedMap(); + verify(mockFetcher).getReadyWindows(); + verify(mockFetcher).prefetchElements(readyWindows); + verify(mockBag1).read(); + verify(mockBag1).clear(); + verify(mockBag2).read(); + verify(mockBag2).clear(); + verify(mockFetcher).releaseBlockedWindows(readyWindows); + } + + @Test + public void testHandleFinishBundle() { + // Given + doNothing().when(mockFetcher).persist(); + + // When + processor.handleFinishBundle(); + + // Then + verify(mockFetcher).persist(); + } + + @Test + public void testHandleProcessElementBlocked() { + // Given + IntervalWindow window = new IntervalWindow(new Instant(0), new Instant(10)); + WindowedValue compressedElement = + WindowedValues.of("e", new Instant(5), Arrays.asList(window), PaneInfo.NO_FIRING); + + when(mockFetcher.storeIfBlocked(any(WindowedValue.class))).thenReturn(true); + + // When + Iterator> unblocked = + processor.handleProcessElement(compressedElement); + + // Then + assertFalse(unblocked.hasNext()); + for (WindowedValue exploded : compressedElement.explodeWindows()) { + verify(mockFetcher).storeIfBlocked(exploded); + } + } + + @Test + public void testHandleProcessElementUnblocked() { + // Given + IntervalWindow window1 = new IntervalWindow(new Instant(0), new Instant(10)); + IntervalWindow window2 = new IntervalWindow(new Instant(10), new Instant(20)); + WindowedValue compressedElement = + WindowedValues.of("e", new Instant(5), Arrays.asList(window1, window2), PaneInfo.NO_FIRING); + + when(mockFetcher.storeIfBlocked(any(WindowedValue.class))).thenReturn(false); + + // When + Iterator> unblocked = + processor.handleProcessElement(compressedElement); + // Then + assertThat( + Lists.newArrayList(unblocked), + containsInAnyOrder( + Iterables.toArray(compressedElement.explodeWindows(), WindowedValue.class))); + for (WindowedValue exploded : compressedElement.explodeWindows()) { + verify(mockFetcher).storeIfBlocked(exploded); + } + } + + @Test + public void testHandleProcessTimerSuccess() { + // Given + TimerData mockTimer = mock(TimerData.class); + when(mockFetcher.storeIfBlocked(mockTimer)).thenReturn(false); + + // When + processor.handleProcessTimer(mockTimer); + + // Then + verify(mockFetcher).storeIfBlocked(mockTimer); + } + + @Test + public void testHandleProcessTimerThrowsPreconditionFail() { + // Given + TimerData mockTimer = mock(TimerData.class); + when(mockFetcher.storeIfBlocked(mockTimer)).thenReturn(true); + + // When & Then + assertThrows( + IllegalStateException.class, + () -> { + processor.handleProcessTimer(mockTimer); + }); + verify(mockFetcher).storeIfBlocked(mockTimer); + } +} diff --git a/runners/prism/java/build.gradle b/runners/prism/java/build.gradle index c89974cb6ea5..6ed720edaeeb 100644 --- a/runners/prism/java/build.gradle +++ b/runners/prism/java/build.gradle @@ -162,6 +162,9 @@ def sickbayTests = [ // java.lang.IllegalStateException: java.io.EOFException 'org.apache.beam.sdk.transforms.ViewTest.testSideInputWithNestedIterables', + // Triggers index-out-of-bound error in Prism + 'org.apache.beam.sdk.transforms.ParDoTest$StateTests.testTimerSideInput', + // Missing output due to processing time timer skew. 'org.apache.beam.sdk.transforms.ParDoTest$TimestampTests.testProcessElementSkew', diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnOp.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnOp.java new file mode 100644 index 000000000000..bc87e2460ec4 --- /dev/null +++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnOp.java @@ -0,0 +1,579 @@ +/* + * 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.samza.runtime; + +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.HashMap; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.ServiceLoader; +import java.util.concurrent.CompletionStage; +import java.util.function.Function; +import org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.runners.core.DoFnRunner; +import org.apache.beam.runners.core.PushbackSideInputDoFnRunner; +import org.apache.beam.runners.core.SideInputHandler; +import org.apache.beam.runners.core.SimplePushbackSideInputDoFnRunner; +import org.apache.beam.runners.core.StateNamespace; +import org.apache.beam.runners.core.StateNamespaces; +import org.apache.beam.runners.core.TimerInternals; +import org.apache.beam.runners.fnexecution.control.ExecutableStageContext; +import org.apache.beam.runners.fnexecution.control.StageBundleFactory; +import org.apache.beam.runners.fnexecution.provisioning.JobInfo; +import org.apache.beam.runners.samza.SamzaExecutionContext; +import org.apache.beam.runners.samza.SamzaPipelineOptions; +import org.apache.beam.runners.samza.util.DoFnUtils; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.DoFnSchemaInformation; +import org.apache.beam.sdk.transforms.join.RawUnionValue; +import org.apache.beam.sdk.transforms.reflect.DoFnInvoker; +import org.apache.beam.sdk.transforms.reflect.DoFnInvokers; +import org.apache.beam.sdk.transforms.reflect.DoFnSignature; +import org.apache.beam.sdk.transforms.reflect.DoFnSignatures; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.WindowedValueMultiReceiver; +import org.apache.beam.sdk.util.construction.graph.ExecutableStage; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionView; +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.Iterators; +import org.apache.samza.config.Config; +import org.apache.samza.context.Context; +import org.apache.samza.operators.Scheduler; +import org.joda.time.Instant; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** Samza operator for {@link DoFn}. */ +@SuppressWarnings({ + "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) + "nullness" // TODO(https://github.com/apache/beam/issues/20497) +}) +public class DoFnOp implements Op { + private static final Logger LOG = LoggerFactory.getLogger(DoFnOp.class); + + private final TupleTag mainOutputTag; + private final DoFn doFn; + private final Coder keyCoder; + private final Collection> sideInputs; + private final List> sideOutputTags; + private final WindowingStrategy windowingStrategy; + private final OutputManagerFactory outputManagerFactory; + // NOTE: we use HashMap here to guarantee Serializability + // Mapping from view id to a view + private final HashMap> idToViewMap; + private final String transformFullName; + private final String transformId; + private final Coder inputCoder; + private final Coder> windowedValueCoder; + private final HashMap, Coder> outputCoders; + private final PCollection.IsBounded isBounded; + private final String bundleCheckTimerId; + private final String bundleStateId; + + // portable api related + private final boolean isPortable; + private final RunnerApi.ExecutableStagePayload stagePayload; + private final JobInfo jobInfo; + private final HashMap> idToTupleTagMap; + + private transient SamzaTimerInternalsFactory timerInternalsFactory; + private transient DoFnRunner fnRunner; + private transient PushbackSideInputDoFnRunner pushbackFnRunner; + private transient SideInputHandler sideInputHandler; + private transient DoFnInvoker doFnInvoker; + private transient SamzaPipelineOptions samzaPipelineOptions; + + // This is derivable from pushbackValues which is persisted to a store. + // TODO: eagerly initialize the hold in init + @edu.umd.cs.findbugs.annotations.SuppressWarnings( + justification = "No bug", + value = "SE_TRANSIENT_FIELD_NOT_RESTORED") + private transient Instant pushbackWatermarkHold; + + // TODO: add this to checkpointable state + private transient Instant inputWatermark; + private transient BundleManager bundleManager; + private transient Instant sideInputWatermark; + private transient List> pushbackValues; + private transient ExecutableStageContext stageContext; + private transient StageBundleFactory stageBundleFactory; + private transient boolean bundleDisabled; + + private final DoFnSchemaInformation doFnSchemaInformation; + private final Map> sideInputMapping; + private final Map stateIdToStoreMapping; + + public DoFnOp( + TupleTag mainOutputTag, + DoFn doFn, + Coder keyCoder, + Coder inputCoder, + Coder> windowedValueCoder, + Map, Coder> outputCoders, + Collection> sideInputs, + List> sideOutputTags, + WindowingStrategy windowingStrategy, + Map> idToViewMap, + OutputManagerFactory outputManagerFactory, + String transformFullName, + String transformId, + PCollection.IsBounded isBounded, + boolean isPortable, + RunnerApi.ExecutableStagePayload stagePayload, + JobInfo jobInfo, + Map> idToTupleTagMap, + DoFnSchemaInformation doFnSchemaInformation, + Map> sideInputMapping, + Map stateIdToStoreMapping) { + this.mainOutputTag = mainOutputTag; + this.doFn = doFn; + this.sideInputs = sideInputs; + this.sideOutputTags = sideOutputTags; + this.inputCoder = inputCoder; + this.windowedValueCoder = windowedValueCoder; + this.outputCoders = new HashMap<>(outputCoders); + this.windowingStrategy = windowingStrategy; + this.idToViewMap = new HashMap<>(idToViewMap); + this.outputManagerFactory = outputManagerFactory; + this.transformFullName = transformFullName; + this.transformId = transformId; + this.keyCoder = keyCoder; + this.isBounded = isBounded; + this.isPortable = isPortable; + this.stagePayload = stagePayload; + this.jobInfo = jobInfo; + this.idToTupleTagMap = new HashMap<>(idToTupleTagMap); + this.bundleCheckTimerId = "_samza_bundle_check_" + transformId; + this.bundleStateId = "_samza_bundle_" + transformId; + this.doFnSchemaInformation = doFnSchemaInformation; + this.sideInputMapping = sideInputMapping; + this.stateIdToStoreMapping = stateIdToStoreMapping; + } + + @Override + @SuppressWarnings("unchecked") + public void open( + Config config, + Context context, + Scheduler> timerRegistry, + OpEmitter emitter) { + this.inputWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE; + this.sideInputWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE; + this.pushbackWatermarkHold = BoundedWindow.TIMESTAMP_MAX_VALUE; + + final DoFnSignature signature = DoFnSignatures.getSignature(doFn.getClass()); + final SamzaExecutionContext samzaExecutionContext = + (SamzaExecutionContext) context.getApplicationContainerContext(); + this.samzaPipelineOptions = samzaExecutionContext.getPipelineOptions(); + this.bundleDisabled = samzaPipelineOptions.getMaxBundleSize() <= 1; + + final String stateId = "pardo-" + transformId; + final SamzaStoreStateInternals.Factory nonKeyedStateInternalsFactory = + SamzaStoreStateInternals.createNonKeyedStateInternalsFactory( + stateId, context.getTaskContext(), samzaPipelineOptions); + final FutureCollector outputFutureCollector = createFutureCollector(); + + this.bundleManager = + isPortable + ? new PortableBundleManager<>( + createBundleProgressListener(), + samzaPipelineOptions.getMaxBundleSize(), + samzaPipelineOptions.getMaxBundleTimeMs(), + timerRegistry, + bundleCheckTimerId) + : new ClassicBundleManager<>( + createBundleProgressListener(), + outputFutureCollector, + samzaPipelineOptions.getMaxBundleSize(), + samzaPipelineOptions.getMaxBundleTimeMs(), + timerRegistry, + bundleCheckTimerId); + + this.timerInternalsFactory = + SamzaTimerInternalsFactory.createTimerInternalFactory( + keyCoder, + (Scheduler) timerRegistry, + getTimerStateId(signature), + nonKeyedStateInternalsFactory, + windowingStrategy, + isBounded, + samzaPipelineOptions); + + this.sideInputHandler = + new SideInputHandler(sideInputs, nonKeyedStateInternalsFactory.stateInternalsForKey(null)); + + if (isPortable) { + final ExecutableStage executableStage = ExecutableStage.fromPayload(stagePayload); + stageContext = SamzaExecutableStageContextFactory.getInstance().get(jobInfo); + stageBundleFactory = stageContext.getStageBundleFactory(executableStage); + this.fnRunner = + SamzaDoFnRunners.createPortable( + transformId, + DoFnUtils.toStepName(executableStage), + bundleStateId, + windowedValueCoder, + executableStage, + sideInputMapping, + sideInputHandler, + nonKeyedStateInternalsFactory, + timerInternalsFactory, + samzaPipelineOptions, + outputManagerFactory.create(emitter, outputFutureCollector), + stageBundleFactory, + samzaExecutionContext, + mainOutputTag, + idToTupleTagMap, + context, + transformFullName); + } else { + this.fnRunner = + SamzaDoFnRunners.create( + samzaPipelineOptions, + doFn, + windowingStrategy, + transformFullName, + stateId, + context, + mainOutputTag, + sideInputHandler, + timerInternalsFactory, + keyCoder, + outputManagerFactory.create(emitter, outputFutureCollector), + inputCoder, + sideOutputTags, + outputCoders, + doFnSchemaInformation, + (Map>) sideInputMapping, + stateIdToStoreMapping, + emitter, + outputFutureCollector); + } + + this.pushbackFnRunner = + SimplePushbackSideInputDoFnRunner.create(fnRunner, sideInputs, sideInputHandler); + this.pushbackValues = new ArrayList<>(); + + final Iterator invokerReg = + ServiceLoader.load(SamzaDoFnInvokerRegistrar.class).iterator(); + if (!invokerReg.hasNext()) { + // use the default invoker here + doFnInvoker = DoFnInvokers.tryInvokeSetupFor(doFn, samzaPipelineOptions); + } else { + doFnInvoker = + Iterators.getOnlyElement(invokerReg).invokerSetupFor(doFn, samzaPipelineOptions, context); + } + } + + FutureCollector createFutureCollector() { + return new FutureCollectorImpl<>(); + } + + private String getTimerStateId(DoFnSignature signature) { + final StringBuilder builder = new StringBuilder("timer"); + if (signature.usesTimers()) { + signature.timerDeclarations().keySet().forEach(builder::append); + } + return builder.toString(); + } + + @Override + public void processElement(WindowedValue inputElement, OpEmitter emitter) { + try { + bundleManager.tryStartBundle(); + final Iterable> rejectedValues = + pushbackFnRunner.processElementInReadyWindows(inputElement); + for (WindowedValue rejectedValue : rejectedValues) { + if (rejectedValue.getTimestamp().compareTo(pushbackWatermarkHold) < 0) { + pushbackWatermarkHold = rejectedValue.getTimestamp(); + } + pushbackValues.add(rejectedValue); + } + + bundleManager.tryFinishBundle(emitter); + } catch (Throwable t) { + LOG.error("Encountered error during process element", t); + bundleManager.signalFailure(t); + throw t; + } + } + + private void doProcessWatermark(Instant watermark, OpEmitter emitter) { + this.inputWatermark = watermark; + + if (sideInputWatermark.isEqual(BoundedWindow.TIMESTAMP_MAX_VALUE)) { + // this means we will never see any more side input + emitAllPushbackValues(); + } + + final Instant actualInputWatermark = + pushbackWatermarkHold.isBefore(inputWatermark) ? pushbackWatermarkHold : inputWatermark; + + timerInternalsFactory.setInputWatermark(actualInputWatermark); + + Collection> readyTimers = timerInternalsFactory.removeReadyTimers(); + if (!readyTimers.isEmpty()) { + pushbackFnRunner.startBundle(); + for (KeyedTimerData keyedTimerData : readyTimers) { + fireTimer(keyedTimerData); + } + pushbackFnRunner.finishBundle(); + } + + if (timerInternalsFactory.getOutputWatermark() == null + || timerInternalsFactory.getOutputWatermark().isBefore(actualInputWatermark)) { + timerInternalsFactory.setOutputWatermark(actualInputWatermark); + emitter.emitWatermark(timerInternalsFactory.getOutputWatermark()); + } + } + + @Override + public void processWatermark(Instant watermark, OpEmitter emitter) { + bundleManager.processWatermark(watermark, emitter); + } + + @Override + public void processSideInput( + String id, WindowedValue> elements, OpEmitter emitter) { + checkState( + bundleDisabled, "Side input not supported in bundling mode. Please disable bundling."); + @SuppressWarnings("unchecked") + final WindowedValue> retypedElements = (WindowedValue>) elements; + + final PCollectionView view = idToViewMap.get(id); + if (view == null) { + throw new IllegalArgumentException("No mapping of id " + id + " to view."); + } + + sideInputHandler.addSideInputValue(view, retypedElements); + + final List> previousPushbackValues = new ArrayList<>(pushbackValues); + pushbackWatermarkHold = BoundedWindow.TIMESTAMP_MAX_VALUE; + pushbackValues.clear(); + + for (final WindowedValue value : previousPushbackValues) { + processElement(value, emitter); + } + + // We may be able to advance the output watermark since we may have played some pushed back + // events. + processWatermark(this.inputWatermark, emitter); + } + + @Override + public void processSideInputWatermark(Instant watermark, OpEmitter emitter) { + checkState( + bundleDisabled, "Side input not supported in bundling mode. Please disable bundling."); + sideInputWatermark = watermark; + + if (sideInputWatermark.isEqual(BoundedWindow.TIMESTAMP_MAX_VALUE)) { + // this means we will never see any more side input + processWatermark(this.inputWatermark, emitter); + } + } + + @Override + @SuppressWarnings("unchecked") + public void processTimer(KeyedTimerData keyedTimerData, OpEmitter emitter) { + // this is internal timer in processing time to check whether a bundle should be closed + if (bundleCheckTimerId.equals(keyedTimerData.getTimerData().getTimerId())) { + bundleManager.processTimer(keyedTimerData, emitter); + return; + } + + pushbackFnRunner.startBundle(); + fireTimer(keyedTimerData); + pushbackFnRunner.finishBundle(); + + this.timerInternalsFactory.removeProcessingTimer((KeyedTimerData) keyedTimerData); + } + + @Override + public void close() { + doFnInvoker.invokeTeardown(); + try (AutoCloseable factory = stageBundleFactory; + AutoCloseable context = stageContext) { + // do nothing + } catch (Exception e) { + LOG.error("Failed to close stage bundle factory", e); + } + } + + private void fireTimer(KeyedTimerData keyedTimerData) { + final TimerInternals.TimerData timer = keyedTimerData.getTimerData(); + LOG.debug("Firing timer {}", timer); + + final StateNamespace namespace = timer.getNamespace(); + // NOTE: not sure why this is safe, but DoFnOperator makes this assumption + final BoundedWindow window = ((StateNamespaces.WindowNamespace) namespace).getWindow(); + + fnRunner.onTimer( + timer.getTimerId(), + timer.getTimerFamilyId(), + keyedTimerData.getKey(), + window, + timer.getTimestamp(), + timer.getOutputTimestamp(), + timer.getDomain(), + timer.causedByDrain()); + } + + // todo: should this go through bundle manager to start and finish the bundle? + private void emitAllPushbackValues() { + if (!pushbackValues.isEmpty()) { + pushbackFnRunner.startBundle(); + + final List> previousPushbackValues = new ArrayList<>(pushbackValues); + pushbackWatermarkHold = BoundedWindow.TIMESTAMP_MAX_VALUE; + pushbackValues.clear(); + + for (final WindowedValue value : previousPushbackValues) { + fnRunner.processElement(value); + } + + pushbackFnRunner.finishBundle(); + } + } + + private BundleManager.BundleProgressListener createBundleProgressListener() { + return new BundleManager.BundleProgressListener() { + @Override + public void onBundleStarted() { + pushbackFnRunner.startBundle(); + } + + @Override + public void onBundleFinished(OpEmitter emitter) { + pushbackFnRunner.finishBundle(); + } + + @Override + public void onWatermark(Instant watermark, OpEmitter emitter) { + doProcessWatermark(watermark, emitter); + } + }; + } + + static CompletionStage> createOutputFuture( + WindowedValue windowedValue, + CompletionStage valueFuture, + Function valueMapper) { + return valueFuture.thenApply( + res -> + WindowedValues.of( + valueMapper.apply(res), + windowedValue.getTimestamp(), + windowedValue.getWindows(), + windowedValue.getPaneInfo())); + } + + /** + * Factory class to create an {@link org.apache.beam.sdk.util.WindowedValueMultiReceiver} that + * emits values to the main output only, which is a single {@link + * org.apache.beam.sdk.values.PCollection}. + * + * @param type of the output element. + */ + public static class SingleOutputManagerFactory implements OutputManagerFactory { + @Override + public WindowedValueMultiReceiver create(OpEmitter emitter) { + return createOutputManager(emitter, null); + } + + @Override + public WindowedValueMultiReceiver create( + OpEmitter emitter, FutureCollector collector) { + return createOutputManager(emitter, collector); + } + + private WindowedValueMultiReceiver createOutputManager( + OpEmitter emitter, FutureCollector collector) { + return new WindowedValueMultiReceiver() { + @Override + @SuppressWarnings("unchecked") + public void output(TupleTag tupleTag, WindowedValue windowedValue) { + // With only one input we know that T is of type OutT. + if (windowedValue.getValue() instanceof CompletionStage) { + CompletionStage valueFuture = (CompletionStage) windowedValue.getValue(); + if (collector != null) { + collector.add(createOutputFuture(windowedValue, valueFuture, value -> (OutT) value)); + } + } else { + final WindowedValue retypedWindowedValue = (WindowedValue) windowedValue; + emitter.emitElement(retypedWindowedValue); + } + } + }; + } + } + + /** + * Factory class to create an {@link org.apache.beam.runners.core.WindowedValueMultiReceiver} that + * emits values to the main output as well as the side outputs via union type {@link + * RawUnionValue}. + */ + public static class MultiOutputManagerFactory implements OutputManagerFactory { + private final Map, Integer> tagToIndexMap; + + public MultiOutputManagerFactory(Map, Integer> tagToIndexMap) { + this.tagToIndexMap = tagToIndexMap; + } + + @Override + public WindowedValueMultiReceiver create(OpEmitter emitter) { + return createOutputManager(emitter, null); + } + + @Override + public WindowedValueMultiReceiver create( + OpEmitter emitter, FutureCollector collector) { + return createOutputManager(emitter, collector); + } + + private WindowedValueMultiReceiver createOutputManager( + OpEmitter emitter, FutureCollector collector) { + return new WindowedValueMultiReceiver() { + @Override + @SuppressWarnings("unchecked") + public void output(TupleTag tupleTag, WindowedValue windowedValue) { + final int index = tagToIndexMap.get(tupleTag); + final T rawValue = windowedValue.getValue(); + if (rawValue instanceof CompletionStage) { + CompletionStage valueFuture = (CompletionStage) rawValue; + if (collector != null) { + collector.add( + createOutputFuture( + windowedValue, valueFuture, res -> new RawUnionValue(index, res))); + } + } else { + final RawUnionValue rawUnionValue = new RawUnionValue(index, rawValue); + emitter.emitElement(windowedValue.withValue(rawUnionValue)); + } + } + }; + } + } +} diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/PortableDoFnOp.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/PortableDoFnOp.java new file mode 100644 index 000000000000..468e4b9aa8dc --- /dev/null +++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/PortableDoFnOp.java @@ -0,0 +1,467 @@ +/* + * 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.samza.runtime; + +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.HashMap; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.ServiceLoader; +import org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.runners.core.DoFnRunner; +import org.apache.beam.runners.core.PushbackSideInputDoFnRunner; +import org.apache.beam.runners.core.SideInputHandler; +import org.apache.beam.runners.core.SimplePushbackSideInputDoFnRunner; +import org.apache.beam.runners.core.StateNamespace; +import org.apache.beam.runners.core.StateNamespaces; +import org.apache.beam.runners.core.TimerInternals; +import org.apache.beam.runners.fnexecution.control.ExecutableStageContext; +import org.apache.beam.runners.fnexecution.control.StageBundleFactory; +import org.apache.beam.runners.fnexecution.provisioning.JobInfo; +import org.apache.beam.runners.samza.SamzaExecutionContext; +import org.apache.beam.runners.samza.SamzaPipelineOptions; +import org.apache.beam.runners.samza.util.DoFnUtils; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.DoFnSchemaInformation; +import org.apache.beam.sdk.transforms.reflect.DoFnInvoker; +import org.apache.beam.sdk.transforms.reflect.DoFnInvokers; +import org.apache.beam.sdk.transforms.reflect.DoFnSignature; +import org.apache.beam.sdk.transforms.reflect.DoFnSignatures; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.construction.graph.ExecutableStage; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionView; +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.collect.Iterators; +import org.apache.samza.config.Config; +import org.apache.samza.context.Context; +import org.apache.samza.operators.Scheduler; +import org.joda.time.Instant; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** Samza operator for {@link DoFn}. */ +@SuppressWarnings({ + "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) + "nullness" // TODO(https://github.com/apache/beam/issues/20497) +}) +public class PortableDoFnOp implements Op { + private static final Logger LOG = LoggerFactory.getLogger(PortableDoFnOp.class); + + private final TupleTag mainOutputTag; + private final DoFn doFn; + private final Coder keyCoder; + private final Collection> sideInputs; + private final List> sideOutputTags; + private final WindowingStrategy windowingStrategy; + private final OutputManagerFactory outputManagerFactory; + // NOTE: we use HashMap here to guarantee Serializability + // Mapping from view id to a view + private final HashMap> idToViewMap; + private final String transformFullName; + private final String transformId; + private final Coder inputCoder; + private final Coder> windowedValueCoder; + private final HashMap, Coder> outputCoders; + private final PCollection.IsBounded isBounded; + private final String bundleCheckTimerId; + private final String bundleStateId; + + // portable api related + private final boolean isPortable; + private final RunnerApi.ExecutableStagePayload stagePayload; + private final JobInfo jobInfo; + private final HashMap> idToTupleTagMap; + + private transient SamzaTimerInternalsFactory timerInternalsFactory; + private transient DoFnRunner fnRunner; + private transient PushbackSideInputDoFnRunner pushbackFnRunner; + private transient SideInputHandler sideInputHandler; + private transient DoFnInvoker doFnInvoker; + private transient SamzaPipelineOptions samzaPipelineOptions; + + // This is derivable from pushbackValues which is persisted to a store. + // TODO: eagerly initialize the hold in init + @edu.umd.cs.findbugs.annotations.SuppressWarnings( + justification = "No bug", + value = "SE_TRANSIENT_FIELD_NOT_RESTORED") + private transient Instant pushbackWatermarkHold; + + // TODO: add this to checkpointable state + private transient Instant inputWatermark; + private transient BundleManager bundleManager; + private transient Instant sideInputWatermark; + private transient List> pushbackValues; + private transient ExecutableStageContext stageContext; + private transient StageBundleFactory stageBundleFactory; + private transient boolean bundleDisabled; + + private final DoFnSchemaInformation doFnSchemaInformation; + private final Map> sideInputMapping; + private final Map stateIdToStoreMapping; + + public PortableDoFnOp( + TupleTag mainOutputTag, + DoFn doFn, + Coder keyCoder, + Coder inputCoder, + Coder> windowedValueCoder, + Map, Coder> outputCoders, + Collection> sideInputs, + List> sideOutputTags, + WindowingStrategy windowingStrategy, + Map> idToViewMap, + OutputManagerFactory outputManagerFactory, + String transformFullName, + String transformId, + PCollection.IsBounded isBounded, + boolean isPortable, + RunnerApi.ExecutableStagePayload stagePayload, + JobInfo jobInfo, + Map> idToTupleTagMap, + DoFnSchemaInformation doFnSchemaInformation, + Map> sideInputMapping, + Map stateIdToStoreMapping) { + this.mainOutputTag = mainOutputTag; + this.doFn = doFn; + this.sideInputs = sideInputs; + this.sideOutputTags = sideOutputTags; + this.inputCoder = inputCoder; + this.windowedValueCoder = windowedValueCoder; + this.outputCoders = new HashMap<>(outputCoders); + this.windowingStrategy = windowingStrategy; + this.idToViewMap = new HashMap<>(idToViewMap); + this.outputManagerFactory = outputManagerFactory; + this.transformFullName = transformFullName; + this.transformId = transformId; + this.keyCoder = keyCoder; + this.isBounded = isBounded; + this.isPortable = isPortable; + this.stagePayload = stagePayload; + this.jobInfo = jobInfo; + this.idToTupleTagMap = new HashMap<>(idToTupleTagMap); + this.bundleCheckTimerId = "_samza_bundle_check_" + transformId; + this.bundleStateId = "_samza_bundle_" + transformId; + this.doFnSchemaInformation = doFnSchemaInformation; + this.sideInputMapping = sideInputMapping; + this.stateIdToStoreMapping = stateIdToStoreMapping; + } + + @Override + @SuppressWarnings("unchecked") + public void open( + Config config, + Context context, + Scheduler> timerRegistry, + OpEmitter emitter) { + this.inputWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE; + this.sideInputWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE; + this.pushbackWatermarkHold = BoundedWindow.TIMESTAMP_MAX_VALUE; + + final DoFnSignature signature = DoFnSignatures.getSignature(doFn.getClass()); + final SamzaExecutionContext samzaExecutionContext = + (SamzaExecutionContext) context.getApplicationContainerContext(); + this.samzaPipelineOptions = samzaExecutionContext.getPipelineOptions(); + this.bundleDisabled = samzaPipelineOptions.getMaxBundleSize() <= 1; + + final String stateId = "pardo-" + transformId; + final SamzaStoreStateInternals.Factory nonKeyedStateInternalsFactory = + SamzaStoreStateInternals.createNonKeyedStateInternalsFactory( + stateId, context.getTaskContext(), samzaPipelineOptions); + final FutureCollector outputFutureCollector = createFutureCollector(); + + this.bundleManager = + new ClassicBundleManager<>( + createBundleProgressListener(), + outputFutureCollector, + samzaPipelineOptions.getMaxBundleSize(), + samzaPipelineOptions.getMaxBundleTimeMs(), + timerRegistry, + bundleCheckTimerId); + + this.timerInternalsFactory = + SamzaTimerInternalsFactory.createTimerInternalFactory( + keyCoder, + (Scheduler) timerRegistry, + getTimerStateId(signature), + nonKeyedStateInternalsFactory, + windowingStrategy, + isBounded, + samzaPipelineOptions); + + this.sideInputHandler = + new SideInputHandler(sideInputs, nonKeyedStateInternalsFactory.stateInternalsForKey(null)); + + if (isPortable) { + final ExecutableStage executableStage = ExecutableStage.fromPayload(stagePayload); + stageContext = SamzaExecutableStageContextFactory.getInstance().get(jobInfo); + stageBundleFactory = stageContext.getStageBundleFactory(executableStage); + this.fnRunner = + SamzaDoFnRunners.createPortable( + transformId, + DoFnUtils.toStepName(executableStage), + bundleStateId, + windowedValueCoder, + executableStage, + sideInputMapping, + sideInputHandler, + nonKeyedStateInternalsFactory, + timerInternalsFactory, + samzaPipelineOptions, + outputManagerFactory.create(emitter, outputFutureCollector), + stageBundleFactory, + samzaExecutionContext, + mainOutputTag, + idToTupleTagMap, + context, + transformFullName); + } else { + this.fnRunner = + SamzaDoFnRunners.create( + samzaPipelineOptions, + doFn, + windowingStrategy, + transformFullName, + stateId, + context, + mainOutputTag, + sideInputHandler, + timerInternalsFactory, + keyCoder, + outputManagerFactory.create(emitter, outputFutureCollector), + inputCoder, + sideOutputTags, + outputCoders, + doFnSchemaInformation, + (Map>) sideInputMapping, + stateIdToStoreMapping, + emitter, + outputFutureCollector); + } + + this.pushbackFnRunner = + SimplePushbackSideInputDoFnRunner.create(fnRunner, sideInputs, sideInputHandler); + this.pushbackValues = new ArrayList<>(); + + final Iterator invokerReg = + ServiceLoader.load(SamzaDoFnInvokerRegistrar.class).iterator(); + if (!invokerReg.hasNext()) { + // use the default invoker here + doFnInvoker = DoFnInvokers.tryInvokeSetupFor(doFn, samzaPipelineOptions); + } else { + doFnInvoker = + Iterators.getOnlyElement(invokerReg).invokerSetupFor(doFn, samzaPipelineOptions, context); + } + } + + FutureCollector createFutureCollector() { + return new FutureCollectorImpl<>(); + } + + private String getTimerStateId(DoFnSignature signature) { + final StringBuilder builder = new StringBuilder("timer"); + if (signature.usesTimers()) { + signature.timerDeclarations().keySet().forEach(builder::append); + } + return builder.toString(); + } + + @Override + public void processElement(WindowedValue inputElement, OpEmitter emitter) { + try { + bundleManager.tryStartBundle(); + final Iterable> rejectedValues = + pushbackFnRunner.processElementInReadyWindows(inputElement); + for (WindowedValue rejectedValue : rejectedValues) { + if (rejectedValue.getTimestamp().compareTo(pushbackWatermarkHold) < 0) { + pushbackWatermarkHold = rejectedValue.getTimestamp(); + } + pushbackValues.add(rejectedValue); + } + + bundleManager.tryFinishBundle(emitter); + } catch (Throwable t) { + LOG.error("Encountered error during process element", t); + bundleManager.signalFailure(t); + throw t; + } + } + + private void doProcessWatermark(Instant watermark, OpEmitter emitter) { + this.inputWatermark = watermark; + + if (sideInputWatermark.isEqual(BoundedWindow.TIMESTAMP_MAX_VALUE)) { + // this means we will never see any more side input + emitAllPushbackValues(); + } + + final Instant actualInputWatermark = + pushbackWatermarkHold.isBefore(inputWatermark) ? pushbackWatermarkHold : inputWatermark; + + timerInternalsFactory.setInputWatermark(actualInputWatermark); + + Collection> readyTimers = timerInternalsFactory.removeReadyTimers(); + if (!readyTimers.isEmpty()) { + pushbackFnRunner.startBundle(); + for (KeyedTimerData keyedTimerData : readyTimers) { + fireTimer(keyedTimerData); + } + pushbackFnRunner.finishBundle(); + } + + if (timerInternalsFactory.getOutputWatermark() == null + || timerInternalsFactory.getOutputWatermark().isBefore(actualInputWatermark)) { + timerInternalsFactory.setOutputWatermark(actualInputWatermark); + emitter.emitWatermark(timerInternalsFactory.getOutputWatermark()); + } + } + + @Override + public void processWatermark(Instant watermark, OpEmitter emitter) { + bundleManager.processWatermark(watermark, emitter); + } + + @Override + public void processSideInput( + String id, WindowedValue> elements, OpEmitter emitter) { + checkState( + bundleDisabled, "Side input not supported in bundling mode. Please disable bundling."); + @SuppressWarnings("unchecked") + final WindowedValue> retypedElements = (WindowedValue>) elements; + + final PCollectionView view = idToViewMap.get(id); + if (view == null) { + throw new IllegalArgumentException("No mapping of id " + id + " to view."); + } + + sideInputHandler.addSideInputValue(view, retypedElements); + + final List> previousPushbackValues = new ArrayList<>(pushbackValues); + pushbackWatermarkHold = BoundedWindow.TIMESTAMP_MAX_VALUE; + pushbackValues.clear(); + + for (final WindowedValue value : previousPushbackValues) { + processElement(value, emitter); + } + + // We may be able to advance the output watermark since we may have played some pushed back + // events. + processWatermark(this.inputWatermark, emitter); + } + + @Override + public void processSideInputWatermark(Instant watermark, OpEmitter emitter) { + checkState( + bundleDisabled, "Side input not supported in bundling mode. Please disable bundling."); + sideInputWatermark = watermark; + + if (sideInputWatermark.isEqual(BoundedWindow.TIMESTAMP_MAX_VALUE)) { + // this means we will never see any more side input + processWatermark(this.inputWatermark, emitter); + } + } + + @Override + @SuppressWarnings("unchecked") + public void processTimer(KeyedTimerData keyedTimerData, OpEmitter emitter) { + // this is internal timer in processing time to check whether a bundle should be closed + if (bundleCheckTimerId.equals(keyedTimerData.getTimerData().getTimerId())) { + bundleManager.processTimer(keyedTimerData, emitter); + return; + } + + pushbackFnRunner.startBundle(); + fireTimer(keyedTimerData); + pushbackFnRunner.finishBundle(); + + this.timerInternalsFactory.removeProcessingTimer((KeyedTimerData) keyedTimerData); + } + + @Override + public void close() { + doFnInvoker.invokeTeardown(); + try (AutoCloseable factory = stageBundleFactory; + AutoCloseable context = stageContext) { + // do nothing + } catch (Exception e) { + LOG.error("Failed to close stage bundle factory", e); + } + } + + private void fireTimer(KeyedTimerData keyedTimerData) { + final TimerInternals.TimerData timer = keyedTimerData.getTimerData(); + LOG.debug("Firing timer {}", timer); + + final StateNamespace namespace = timer.getNamespace(); + // NOTE: not sure why this is safe, but DoFnOperator makes this assumption + final BoundedWindow window = ((StateNamespaces.WindowNamespace) namespace).getWindow(); + + fnRunner.onTimer( + timer.getTimerId(), + timer.getTimerFamilyId(), + keyedTimerData.getKey(), + window, + timer.getTimestamp(), + timer.getOutputTimestamp(), + timer.getDomain(), + timer.causedByDrain()); + } + + // todo: should this go through bundle manager to start and finish the bundle? + private void emitAllPushbackValues() { + if (!pushbackValues.isEmpty()) { + pushbackFnRunner.startBundle(); + + final List> previousPushbackValues = new ArrayList<>(pushbackValues); + pushbackWatermarkHold = BoundedWindow.TIMESTAMP_MAX_VALUE; + pushbackValues.clear(); + + for (final WindowedValue value : previousPushbackValues) { + fnRunner.processElement(value); + } + + pushbackFnRunner.finishBundle(); + } + } + + private BundleManager.BundleProgressListener createBundleProgressListener() { + return new BundleManager.BundleProgressListener() { + @Override + public void onBundleStarted() { + pushbackFnRunner.startBundle(); + } + + @Override + public void onBundleFinished(OpEmitter emitter) { + pushbackFnRunner.finishBundle(); + } + + @Override + public void onWatermark(Instant watermark, OpEmitter emitter) { + doProcessWatermark(watermark, emitter); + } + }; + } +} diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/UsesSideInputsInTimer.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/UsesSideInputsInTimer.java new file mode 100644 index 000000000000..8320c1451d17 --- /dev/null +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/UsesSideInputsInTimer.java @@ -0,0 +1,27 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.testing; + +import org.apache.beam.sdk.annotations.Internal; + +/** + * Category tag for validation tests which use sideinputs in OnTimer and OnWindowExpiration. Tests + * tagged with {@link UsesSideInputsInTimer} should be run for runners which support sideinputs. + */ +@Internal +public class UsesSideInputsInTimer {} diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFn.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFn.java index a366ded4fe2d..bfd04908b990 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFn.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFn.java @@ -359,6 +359,14 @@ public abstract class OnTimerContext extends WindowedContext { /** Returns the time domain of the current timer. */ public abstract TimeDomain timeDomain(); + /** + * Returns the value of the side input. + * + * @throws IllegalArgumentException if this is not a side input + */ + @Pure + public abstract T sideInput(PCollectionView view); + @Pure public abstract org.apache.beam.sdk.values.CausedByDrain causedByDrain(); } @@ -368,6 +376,14 @@ public abstract class OnWindowExpirationContext extends WindowedContext { /** Returns the window in which the window expiration is firing. */ @Pure public abstract BoundedWindow window(); + + /** + * Returns the value of the side input. + * + * @throws IllegalArgumentException if this is not a side input + */ + @Pure + public abstract T sideInput(PCollectionView view); } /** diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java index 0bd2c1c888f0..2983fc94021c 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java @@ -197,7 +197,8 @@ private DoFnSignatures() {} Parameter.TimerIdParameter.class, Parameter.FireTimestampParameter.class, Parameter.CausedByDrainParameter.class, - Parameter.KeyParameter.class); + Parameter.KeyParameter.class, + Parameter.SideInputParameter.class); private static final ImmutableList> ALLOWED_ON_TIMER_FAMILY_PARAMETERS = @@ -215,7 +216,8 @@ private DoFnSignatures() {} Parameter.TimerIdParameter.class, Parameter.FireTimestampParameter.class, Parameter.CausedByDrainParameter.class, - Parameter.KeyParameter.class); + Parameter.KeyParameter.class, + Parameter.SideInputParameter.class); private static final Collection> ALLOWED_ON_WINDOW_EXPIRATION_PARAMETERS = @@ -226,7 +228,8 @@ private DoFnSignatures() {} Parameter.TaggedOutputReceiverParameter.class, Parameter.StateParameter.class, Parameter.TimestampParameter.class, - Parameter.KeyParameter.class); + Parameter.KeyParameter.class, + Parameter.SideInputParameter.class); private static final Collection> ALLOWED_GET_INITIAL_RESTRICTION_PARAMETERS = diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java index 8a273127b4fc..0c984d01c8f0 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java @@ -114,6 +114,7 @@ import org.apache.beam.sdk.testing.UsesRequiresTimeSortedInput; import org.apache.beam.sdk.testing.UsesSetState; import org.apache.beam.sdk.testing.UsesSideInputs; +import org.apache.beam.sdk.testing.UsesSideInputsInTimer; import org.apache.beam.sdk.testing.UsesSideInputsWithDifferentCoders; import org.apache.beam.sdk.testing.UsesStatefulParDo; import org.apache.beam.sdk.testing.UsesStrictTimerOrdering; @@ -3678,6 +3679,154 @@ public void processElement( pipeline.run(); } + @Test + @Category({ + ValidatesRunner.class, + UsesStatefulParDo.class, + UsesSideInputs.class, + UsesSideInputsInTimer.class, + UsesTestStream.class, + UsesTimersInParDo.class, + UsesTriggeredSideInputs.class, + UsesOnWindowExpiration.class + }) + public void testTimerSideInput() { + // SideInput tag id + final String sideInputTag1 = "tag1"; + + final PCollectionView sideInput = + pipeline + .apply("CreateSideInput1", Create.of(2)) + .apply("ViewSideInput1", View.asSingleton()); + + DoFn, KV> doFn = + new DoFn, KV>() { + @TimerId("timer") + private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.EVENT_TIME); + + @StateId("foo") + private final StateSpec> stateSpec = StateSpecs.value(); + + @ProcessElement + public void process(@Timestamp Instant ts, @TimerId("timer") Timer timer) { + timer.align(Duration.standardSeconds(10)).setRelative(); + } + + @OnTimer("timer") + public void onTimer( + OutputReceiver> o, + @DoFn.SideInput(sideInputTag1) Integer sideInput, + @Key Integer key) { + o.output(KV.of(key, sideInput)); + } + + @OnWindowExpiration + public void onWindowExpiration( + @DoFn.SideInput(sideInputTag1) Integer sideInput, + OutputReceiver> o, + @Key Integer key) { + o.output(KV.of(key, sideInput)); + } + }; + + final int numTestElements = 10; + final Instant now = new Instant(0); + TestStream.Builder> builder = + TestStream.create(KvCoder.of(VarIntCoder.of(), VarIntCoder.of())) + .advanceWatermarkTo(new Instant(0)); + + for (int i = 0; i < numTestElements; i++) { + builder = + builder.addElements( + TimestampedValue.of(KV.of(i % 2, i), now.plus(Duration.millis(i * 1000)))); + if ((i + 1) % 10 == 0) { + builder = builder.advanceWatermarkTo(now.plus(Duration.millis((i + 1) * 1000))); + } + } + List> expected = + IntStream.rangeClosed(0, 1) + .boxed() + .flatMap(i -> ImmutableList.of(KV.of(i, 2), KV.of(i, 2)).stream()) + .collect(Collectors.toList()); + + PCollection> output = + pipeline + .apply(builder.advanceWatermarkToInfinity()) + .apply(ParDo.of(doFn).withSideInput(sideInputTag1, sideInput)); + PAssert.that(output).containsInAnyOrder(expected); + pipeline.run(); + } + + @Test + @Category({ + ValidatesRunner.class, + UsesStatefulParDo.class, + UsesSideInputs.class, + UsesSideInputsInTimer.class, + UsesTimersInParDo.class, + UsesTriggeredSideInputs.class + }) + public void testSideInputNotReadyTimer() { + final String sideInputTag = "tag1"; + + // Create a side input that is delayed by 5 seconds using Thread.sleep + DoFn, String> delayFn = + new DoFn, String>() { + @ProcessElement + public void process(OutputReceiver o) throws InterruptedException { + Thread.sleep(java.time.Duration.ofSeconds(15).toMillis()); + o.output("side-value"); + } + }; + + PCollectionView sideInput = + pipeline + .apply("CreateSideSource", Create.of(KV.of("dummyKey", ""))) + .apply("DelaySideInput", ParDo.of(delayFn)) + .apply(View.asSingleton()); + + // Main input in global window + DoFn, String> fn = + new DoFn, String>() { + @TimerId("timer") + private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.EVENT_TIME); + + @StateId("dummy") + private final StateSpec> dummy = StateSpecs.value(); + + @ProcessElement + public void process( + @Timestamp Instant ts, + @TimerId("timer") Timer timer, + @DoFn.SideInput(sideInputTag) String sideInputValue, + OutputReceiver o) { + // Set timer to fire at current timestamp + 1 millis + timer.offset(Duration.millis(1)).setRelative(); + o.output(sideInputValue); + } + + @OnTimer("timer") + public void onTimer( + OutputReceiver o, @DoFn.SideInput(sideInputTag) String sideInputValue) { + o.output(sideInputValue); + } + + @OnWindowExpiration + public void onWindowExpiration( + OutputReceiver o, @DoFn.SideInput(sideInputTag) String sideInputValue) { + o.output(sideInputValue); + } + }; + + PCollection output = + pipeline + .apply("CreateMainKV", Create.of(KV.of("key", "main-elem"))) + .apply(ParDo.of(fn).withSideInput(sideInputTag, sideInput)); + + PAssert.that(output).containsInAnyOrder("side-value", "side-value", "side-value"); + pipeline.run(); + } + @Test @Category({ ValidatesRunner.class, diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java index d92a84ea9ff7..100a392a9ed7 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java @@ -2362,6 +2362,11 @@ public BoundedWindow window() { return currentWindow; } + @Override + public T sideInput(PCollectionView view) { + return stateAccessor.get(view, currentWindow); + } + @Override public OutputBuilder builder(OutputT value) { return WindowedValues.builder() @@ -2489,6 +2494,15 @@ public K key() { return (K) currentTimer.getUserKey(); } + @Override + public @Nullable Object sideInput(String tagId) { + PCollectionView view = sideInputMapping.get(tagId); + if (view == null) { + throw new IllegalArgumentException("Unknown side input: " + tagId); + } + return stateAccessor.get(view, currentWindow); + } + @Override public OutputReceiver outputReceiver(DoFn doFn) { return context; @@ -2649,6 +2663,11 @@ public BoundedWindow window() { return currentWindow; } + @Override + public T sideInput(PCollectionView view) { + return stateAccessor.get(view, currentWindow); + } + @Override public CausedByDrain causedByDrain() { return causedByDrain; @@ -2800,6 +2819,14 @@ public Instant fireTimestamp(DoFn doFn) { return currentTimer.getFireTimestamp(); } + public @Nullable Object sideInput(String tagId) { + PCollectionView view = sideInputMapping.get(tagId); + if (view == null) { + throw new IllegalArgumentException("Unknown side input: " + tagId); + } + return stateAccessor.get(view, currentWindow); + } + @Override public K key() { return (K) currentTimer.getUserKey(); From 216e99264073c60040738fb6379e5f43ac61d9ac Mon Sep 17 00:00:00 2001 From: Reuven Lax Date: Thu, 28 May 2026 13:23:49 -0700 Subject: [PATCH 02/11] remove files --- .../beam/runners/samza/runtime/DoFnOp.java | 579 ------------------ .../runners/samza/runtime/PortableDoFnOp.java | 467 -------------- 2 files changed, 1046 deletions(-) delete mode 100644 runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnOp.java delete mode 100644 runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/PortableDoFnOp.java diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnOp.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnOp.java deleted file mode 100644 index bc87e2460ec4..000000000000 --- a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnOp.java +++ /dev/null @@ -1,579 +0,0 @@ -/* - * 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.samza.runtime; - -import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; - -import java.util.ArrayList; -import java.util.Collection; -import java.util.HashMap; -import java.util.Iterator; -import java.util.List; -import java.util.Map; -import java.util.ServiceLoader; -import java.util.concurrent.CompletionStage; -import java.util.function.Function; -import org.apache.beam.model.pipeline.v1.RunnerApi; -import org.apache.beam.runners.core.DoFnRunner; -import org.apache.beam.runners.core.PushbackSideInputDoFnRunner; -import org.apache.beam.runners.core.SideInputHandler; -import org.apache.beam.runners.core.SimplePushbackSideInputDoFnRunner; -import org.apache.beam.runners.core.StateNamespace; -import org.apache.beam.runners.core.StateNamespaces; -import org.apache.beam.runners.core.TimerInternals; -import org.apache.beam.runners.fnexecution.control.ExecutableStageContext; -import org.apache.beam.runners.fnexecution.control.StageBundleFactory; -import org.apache.beam.runners.fnexecution.provisioning.JobInfo; -import org.apache.beam.runners.samza.SamzaExecutionContext; -import org.apache.beam.runners.samza.SamzaPipelineOptions; -import org.apache.beam.runners.samza.util.DoFnUtils; -import org.apache.beam.sdk.coders.Coder; -import org.apache.beam.sdk.transforms.DoFn; -import org.apache.beam.sdk.transforms.DoFnSchemaInformation; -import org.apache.beam.sdk.transforms.join.RawUnionValue; -import org.apache.beam.sdk.transforms.reflect.DoFnInvoker; -import org.apache.beam.sdk.transforms.reflect.DoFnInvokers; -import org.apache.beam.sdk.transforms.reflect.DoFnSignature; -import org.apache.beam.sdk.transforms.reflect.DoFnSignatures; -import org.apache.beam.sdk.transforms.windowing.BoundedWindow; -import org.apache.beam.sdk.util.WindowedValueMultiReceiver; -import org.apache.beam.sdk.util.construction.graph.ExecutableStage; -import org.apache.beam.sdk.values.PCollection; -import org.apache.beam.sdk.values.PCollectionView; -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.Iterators; -import org.apache.samza.config.Config; -import org.apache.samza.context.Context; -import org.apache.samza.operators.Scheduler; -import org.joda.time.Instant; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -/** Samza operator for {@link DoFn}. */ -@SuppressWarnings({ - "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) - "nullness" // TODO(https://github.com/apache/beam/issues/20497) -}) -public class DoFnOp implements Op { - private static final Logger LOG = LoggerFactory.getLogger(DoFnOp.class); - - private final TupleTag mainOutputTag; - private final DoFn doFn; - private final Coder keyCoder; - private final Collection> sideInputs; - private final List> sideOutputTags; - private final WindowingStrategy windowingStrategy; - private final OutputManagerFactory outputManagerFactory; - // NOTE: we use HashMap here to guarantee Serializability - // Mapping from view id to a view - private final HashMap> idToViewMap; - private final String transformFullName; - private final String transformId; - private final Coder inputCoder; - private final Coder> windowedValueCoder; - private final HashMap, Coder> outputCoders; - private final PCollection.IsBounded isBounded; - private final String bundleCheckTimerId; - private final String bundleStateId; - - // portable api related - private final boolean isPortable; - private final RunnerApi.ExecutableStagePayload stagePayload; - private final JobInfo jobInfo; - private final HashMap> idToTupleTagMap; - - private transient SamzaTimerInternalsFactory timerInternalsFactory; - private transient DoFnRunner fnRunner; - private transient PushbackSideInputDoFnRunner pushbackFnRunner; - private transient SideInputHandler sideInputHandler; - private transient DoFnInvoker doFnInvoker; - private transient SamzaPipelineOptions samzaPipelineOptions; - - // This is derivable from pushbackValues which is persisted to a store. - // TODO: eagerly initialize the hold in init - @edu.umd.cs.findbugs.annotations.SuppressWarnings( - justification = "No bug", - value = "SE_TRANSIENT_FIELD_NOT_RESTORED") - private transient Instant pushbackWatermarkHold; - - // TODO: add this to checkpointable state - private transient Instant inputWatermark; - private transient BundleManager bundleManager; - private transient Instant sideInputWatermark; - private transient List> pushbackValues; - private transient ExecutableStageContext stageContext; - private transient StageBundleFactory stageBundleFactory; - private transient boolean bundleDisabled; - - private final DoFnSchemaInformation doFnSchemaInformation; - private final Map> sideInputMapping; - private final Map stateIdToStoreMapping; - - public DoFnOp( - TupleTag mainOutputTag, - DoFn doFn, - Coder keyCoder, - Coder inputCoder, - Coder> windowedValueCoder, - Map, Coder> outputCoders, - Collection> sideInputs, - List> sideOutputTags, - WindowingStrategy windowingStrategy, - Map> idToViewMap, - OutputManagerFactory outputManagerFactory, - String transformFullName, - String transformId, - PCollection.IsBounded isBounded, - boolean isPortable, - RunnerApi.ExecutableStagePayload stagePayload, - JobInfo jobInfo, - Map> idToTupleTagMap, - DoFnSchemaInformation doFnSchemaInformation, - Map> sideInputMapping, - Map stateIdToStoreMapping) { - this.mainOutputTag = mainOutputTag; - this.doFn = doFn; - this.sideInputs = sideInputs; - this.sideOutputTags = sideOutputTags; - this.inputCoder = inputCoder; - this.windowedValueCoder = windowedValueCoder; - this.outputCoders = new HashMap<>(outputCoders); - this.windowingStrategy = windowingStrategy; - this.idToViewMap = new HashMap<>(idToViewMap); - this.outputManagerFactory = outputManagerFactory; - this.transformFullName = transformFullName; - this.transformId = transformId; - this.keyCoder = keyCoder; - this.isBounded = isBounded; - this.isPortable = isPortable; - this.stagePayload = stagePayload; - this.jobInfo = jobInfo; - this.idToTupleTagMap = new HashMap<>(idToTupleTagMap); - this.bundleCheckTimerId = "_samza_bundle_check_" + transformId; - this.bundleStateId = "_samza_bundle_" + transformId; - this.doFnSchemaInformation = doFnSchemaInformation; - this.sideInputMapping = sideInputMapping; - this.stateIdToStoreMapping = stateIdToStoreMapping; - } - - @Override - @SuppressWarnings("unchecked") - public void open( - Config config, - Context context, - Scheduler> timerRegistry, - OpEmitter emitter) { - this.inputWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE; - this.sideInputWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE; - this.pushbackWatermarkHold = BoundedWindow.TIMESTAMP_MAX_VALUE; - - final DoFnSignature signature = DoFnSignatures.getSignature(doFn.getClass()); - final SamzaExecutionContext samzaExecutionContext = - (SamzaExecutionContext) context.getApplicationContainerContext(); - this.samzaPipelineOptions = samzaExecutionContext.getPipelineOptions(); - this.bundleDisabled = samzaPipelineOptions.getMaxBundleSize() <= 1; - - final String stateId = "pardo-" + transformId; - final SamzaStoreStateInternals.Factory nonKeyedStateInternalsFactory = - SamzaStoreStateInternals.createNonKeyedStateInternalsFactory( - stateId, context.getTaskContext(), samzaPipelineOptions); - final FutureCollector outputFutureCollector = createFutureCollector(); - - this.bundleManager = - isPortable - ? new PortableBundleManager<>( - createBundleProgressListener(), - samzaPipelineOptions.getMaxBundleSize(), - samzaPipelineOptions.getMaxBundleTimeMs(), - timerRegistry, - bundleCheckTimerId) - : new ClassicBundleManager<>( - createBundleProgressListener(), - outputFutureCollector, - samzaPipelineOptions.getMaxBundleSize(), - samzaPipelineOptions.getMaxBundleTimeMs(), - timerRegistry, - bundleCheckTimerId); - - this.timerInternalsFactory = - SamzaTimerInternalsFactory.createTimerInternalFactory( - keyCoder, - (Scheduler) timerRegistry, - getTimerStateId(signature), - nonKeyedStateInternalsFactory, - windowingStrategy, - isBounded, - samzaPipelineOptions); - - this.sideInputHandler = - new SideInputHandler(sideInputs, nonKeyedStateInternalsFactory.stateInternalsForKey(null)); - - if (isPortable) { - final ExecutableStage executableStage = ExecutableStage.fromPayload(stagePayload); - stageContext = SamzaExecutableStageContextFactory.getInstance().get(jobInfo); - stageBundleFactory = stageContext.getStageBundleFactory(executableStage); - this.fnRunner = - SamzaDoFnRunners.createPortable( - transformId, - DoFnUtils.toStepName(executableStage), - bundleStateId, - windowedValueCoder, - executableStage, - sideInputMapping, - sideInputHandler, - nonKeyedStateInternalsFactory, - timerInternalsFactory, - samzaPipelineOptions, - outputManagerFactory.create(emitter, outputFutureCollector), - stageBundleFactory, - samzaExecutionContext, - mainOutputTag, - idToTupleTagMap, - context, - transformFullName); - } else { - this.fnRunner = - SamzaDoFnRunners.create( - samzaPipelineOptions, - doFn, - windowingStrategy, - transformFullName, - stateId, - context, - mainOutputTag, - sideInputHandler, - timerInternalsFactory, - keyCoder, - outputManagerFactory.create(emitter, outputFutureCollector), - inputCoder, - sideOutputTags, - outputCoders, - doFnSchemaInformation, - (Map>) sideInputMapping, - stateIdToStoreMapping, - emitter, - outputFutureCollector); - } - - this.pushbackFnRunner = - SimplePushbackSideInputDoFnRunner.create(fnRunner, sideInputs, sideInputHandler); - this.pushbackValues = new ArrayList<>(); - - final Iterator invokerReg = - ServiceLoader.load(SamzaDoFnInvokerRegistrar.class).iterator(); - if (!invokerReg.hasNext()) { - // use the default invoker here - doFnInvoker = DoFnInvokers.tryInvokeSetupFor(doFn, samzaPipelineOptions); - } else { - doFnInvoker = - Iterators.getOnlyElement(invokerReg).invokerSetupFor(doFn, samzaPipelineOptions, context); - } - } - - FutureCollector createFutureCollector() { - return new FutureCollectorImpl<>(); - } - - private String getTimerStateId(DoFnSignature signature) { - final StringBuilder builder = new StringBuilder("timer"); - if (signature.usesTimers()) { - signature.timerDeclarations().keySet().forEach(builder::append); - } - return builder.toString(); - } - - @Override - public void processElement(WindowedValue inputElement, OpEmitter emitter) { - try { - bundleManager.tryStartBundle(); - final Iterable> rejectedValues = - pushbackFnRunner.processElementInReadyWindows(inputElement); - for (WindowedValue rejectedValue : rejectedValues) { - if (rejectedValue.getTimestamp().compareTo(pushbackWatermarkHold) < 0) { - pushbackWatermarkHold = rejectedValue.getTimestamp(); - } - pushbackValues.add(rejectedValue); - } - - bundleManager.tryFinishBundle(emitter); - } catch (Throwable t) { - LOG.error("Encountered error during process element", t); - bundleManager.signalFailure(t); - throw t; - } - } - - private void doProcessWatermark(Instant watermark, OpEmitter emitter) { - this.inputWatermark = watermark; - - if (sideInputWatermark.isEqual(BoundedWindow.TIMESTAMP_MAX_VALUE)) { - // this means we will never see any more side input - emitAllPushbackValues(); - } - - final Instant actualInputWatermark = - pushbackWatermarkHold.isBefore(inputWatermark) ? pushbackWatermarkHold : inputWatermark; - - timerInternalsFactory.setInputWatermark(actualInputWatermark); - - Collection> readyTimers = timerInternalsFactory.removeReadyTimers(); - if (!readyTimers.isEmpty()) { - pushbackFnRunner.startBundle(); - for (KeyedTimerData keyedTimerData : readyTimers) { - fireTimer(keyedTimerData); - } - pushbackFnRunner.finishBundle(); - } - - if (timerInternalsFactory.getOutputWatermark() == null - || timerInternalsFactory.getOutputWatermark().isBefore(actualInputWatermark)) { - timerInternalsFactory.setOutputWatermark(actualInputWatermark); - emitter.emitWatermark(timerInternalsFactory.getOutputWatermark()); - } - } - - @Override - public void processWatermark(Instant watermark, OpEmitter emitter) { - bundleManager.processWatermark(watermark, emitter); - } - - @Override - public void processSideInput( - String id, WindowedValue> elements, OpEmitter emitter) { - checkState( - bundleDisabled, "Side input not supported in bundling mode. Please disable bundling."); - @SuppressWarnings("unchecked") - final WindowedValue> retypedElements = (WindowedValue>) elements; - - final PCollectionView view = idToViewMap.get(id); - if (view == null) { - throw new IllegalArgumentException("No mapping of id " + id + " to view."); - } - - sideInputHandler.addSideInputValue(view, retypedElements); - - final List> previousPushbackValues = new ArrayList<>(pushbackValues); - pushbackWatermarkHold = BoundedWindow.TIMESTAMP_MAX_VALUE; - pushbackValues.clear(); - - for (final WindowedValue value : previousPushbackValues) { - processElement(value, emitter); - } - - // We may be able to advance the output watermark since we may have played some pushed back - // events. - processWatermark(this.inputWatermark, emitter); - } - - @Override - public void processSideInputWatermark(Instant watermark, OpEmitter emitter) { - checkState( - bundleDisabled, "Side input not supported in bundling mode. Please disable bundling."); - sideInputWatermark = watermark; - - if (sideInputWatermark.isEqual(BoundedWindow.TIMESTAMP_MAX_VALUE)) { - // this means we will never see any more side input - processWatermark(this.inputWatermark, emitter); - } - } - - @Override - @SuppressWarnings("unchecked") - public void processTimer(KeyedTimerData keyedTimerData, OpEmitter emitter) { - // this is internal timer in processing time to check whether a bundle should be closed - if (bundleCheckTimerId.equals(keyedTimerData.getTimerData().getTimerId())) { - bundleManager.processTimer(keyedTimerData, emitter); - return; - } - - pushbackFnRunner.startBundle(); - fireTimer(keyedTimerData); - pushbackFnRunner.finishBundle(); - - this.timerInternalsFactory.removeProcessingTimer((KeyedTimerData) keyedTimerData); - } - - @Override - public void close() { - doFnInvoker.invokeTeardown(); - try (AutoCloseable factory = stageBundleFactory; - AutoCloseable context = stageContext) { - // do nothing - } catch (Exception e) { - LOG.error("Failed to close stage bundle factory", e); - } - } - - private void fireTimer(KeyedTimerData keyedTimerData) { - final TimerInternals.TimerData timer = keyedTimerData.getTimerData(); - LOG.debug("Firing timer {}", timer); - - final StateNamespace namespace = timer.getNamespace(); - // NOTE: not sure why this is safe, but DoFnOperator makes this assumption - final BoundedWindow window = ((StateNamespaces.WindowNamespace) namespace).getWindow(); - - fnRunner.onTimer( - timer.getTimerId(), - timer.getTimerFamilyId(), - keyedTimerData.getKey(), - window, - timer.getTimestamp(), - timer.getOutputTimestamp(), - timer.getDomain(), - timer.causedByDrain()); - } - - // todo: should this go through bundle manager to start and finish the bundle? - private void emitAllPushbackValues() { - if (!pushbackValues.isEmpty()) { - pushbackFnRunner.startBundle(); - - final List> previousPushbackValues = new ArrayList<>(pushbackValues); - pushbackWatermarkHold = BoundedWindow.TIMESTAMP_MAX_VALUE; - pushbackValues.clear(); - - for (final WindowedValue value : previousPushbackValues) { - fnRunner.processElement(value); - } - - pushbackFnRunner.finishBundle(); - } - } - - private BundleManager.BundleProgressListener createBundleProgressListener() { - return new BundleManager.BundleProgressListener() { - @Override - public void onBundleStarted() { - pushbackFnRunner.startBundle(); - } - - @Override - public void onBundleFinished(OpEmitter emitter) { - pushbackFnRunner.finishBundle(); - } - - @Override - public void onWatermark(Instant watermark, OpEmitter emitter) { - doProcessWatermark(watermark, emitter); - } - }; - } - - static CompletionStage> createOutputFuture( - WindowedValue windowedValue, - CompletionStage valueFuture, - Function valueMapper) { - return valueFuture.thenApply( - res -> - WindowedValues.of( - valueMapper.apply(res), - windowedValue.getTimestamp(), - windowedValue.getWindows(), - windowedValue.getPaneInfo())); - } - - /** - * Factory class to create an {@link org.apache.beam.sdk.util.WindowedValueMultiReceiver} that - * emits values to the main output only, which is a single {@link - * org.apache.beam.sdk.values.PCollection}. - * - * @param type of the output element. - */ - public static class SingleOutputManagerFactory implements OutputManagerFactory { - @Override - public WindowedValueMultiReceiver create(OpEmitter emitter) { - return createOutputManager(emitter, null); - } - - @Override - public WindowedValueMultiReceiver create( - OpEmitter emitter, FutureCollector collector) { - return createOutputManager(emitter, collector); - } - - private WindowedValueMultiReceiver createOutputManager( - OpEmitter emitter, FutureCollector collector) { - return new WindowedValueMultiReceiver() { - @Override - @SuppressWarnings("unchecked") - public void output(TupleTag tupleTag, WindowedValue windowedValue) { - // With only one input we know that T is of type OutT. - if (windowedValue.getValue() instanceof CompletionStage) { - CompletionStage valueFuture = (CompletionStage) windowedValue.getValue(); - if (collector != null) { - collector.add(createOutputFuture(windowedValue, valueFuture, value -> (OutT) value)); - } - } else { - final WindowedValue retypedWindowedValue = (WindowedValue) windowedValue; - emitter.emitElement(retypedWindowedValue); - } - } - }; - } - } - - /** - * Factory class to create an {@link org.apache.beam.runners.core.WindowedValueMultiReceiver} that - * emits values to the main output as well as the side outputs via union type {@link - * RawUnionValue}. - */ - public static class MultiOutputManagerFactory implements OutputManagerFactory { - private final Map, Integer> tagToIndexMap; - - public MultiOutputManagerFactory(Map, Integer> tagToIndexMap) { - this.tagToIndexMap = tagToIndexMap; - } - - @Override - public WindowedValueMultiReceiver create(OpEmitter emitter) { - return createOutputManager(emitter, null); - } - - @Override - public WindowedValueMultiReceiver create( - OpEmitter emitter, FutureCollector collector) { - return createOutputManager(emitter, collector); - } - - private WindowedValueMultiReceiver createOutputManager( - OpEmitter emitter, FutureCollector collector) { - return new WindowedValueMultiReceiver() { - @Override - @SuppressWarnings("unchecked") - public void output(TupleTag tupleTag, WindowedValue windowedValue) { - final int index = tagToIndexMap.get(tupleTag); - final T rawValue = windowedValue.getValue(); - if (rawValue instanceof CompletionStage) { - CompletionStage valueFuture = (CompletionStage) rawValue; - if (collector != null) { - collector.add( - createOutputFuture( - windowedValue, valueFuture, res -> new RawUnionValue(index, res))); - } - } else { - final RawUnionValue rawUnionValue = new RawUnionValue(index, rawValue); - emitter.emitElement(windowedValue.withValue(rawUnionValue)); - } - } - }; - } - } -} diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/PortableDoFnOp.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/PortableDoFnOp.java deleted file mode 100644 index 468e4b9aa8dc..000000000000 --- a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/PortableDoFnOp.java +++ /dev/null @@ -1,467 +0,0 @@ -/* - * 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.samza.runtime; - -import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; - -import java.util.ArrayList; -import java.util.Collection; -import java.util.HashMap; -import java.util.Iterator; -import java.util.List; -import java.util.Map; -import java.util.ServiceLoader; -import org.apache.beam.model.pipeline.v1.RunnerApi; -import org.apache.beam.runners.core.DoFnRunner; -import org.apache.beam.runners.core.PushbackSideInputDoFnRunner; -import org.apache.beam.runners.core.SideInputHandler; -import org.apache.beam.runners.core.SimplePushbackSideInputDoFnRunner; -import org.apache.beam.runners.core.StateNamespace; -import org.apache.beam.runners.core.StateNamespaces; -import org.apache.beam.runners.core.TimerInternals; -import org.apache.beam.runners.fnexecution.control.ExecutableStageContext; -import org.apache.beam.runners.fnexecution.control.StageBundleFactory; -import org.apache.beam.runners.fnexecution.provisioning.JobInfo; -import org.apache.beam.runners.samza.SamzaExecutionContext; -import org.apache.beam.runners.samza.SamzaPipelineOptions; -import org.apache.beam.runners.samza.util.DoFnUtils; -import org.apache.beam.sdk.coders.Coder; -import org.apache.beam.sdk.transforms.DoFn; -import org.apache.beam.sdk.transforms.DoFnSchemaInformation; -import org.apache.beam.sdk.transforms.reflect.DoFnInvoker; -import org.apache.beam.sdk.transforms.reflect.DoFnInvokers; -import org.apache.beam.sdk.transforms.reflect.DoFnSignature; -import org.apache.beam.sdk.transforms.reflect.DoFnSignatures; -import org.apache.beam.sdk.transforms.windowing.BoundedWindow; -import org.apache.beam.sdk.util.construction.graph.ExecutableStage; -import org.apache.beam.sdk.values.PCollection; -import org.apache.beam.sdk.values.PCollectionView; -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.collect.Iterators; -import org.apache.samza.config.Config; -import org.apache.samza.context.Context; -import org.apache.samza.operators.Scheduler; -import org.joda.time.Instant; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -/** Samza operator for {@link DoFn}. */ -@SuppressWarnings({ - "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) - "nullness" // TODO(https://github.com/apache/beam/issues/20497) -}) -public class PortableDoFnOp implements Op { - private static final Logger LOG = LoggerFactory.getLogger(PortableDoFnOp.class); - - private final TupleTag mainOutputTag; - private final DoFn doFn; - private final Coder keyCoder; - private final Collection> sideInputs; - private final List> sideOutputTags; - private final WindowingStrategy windowingStrategy; - private final OutputManagerFactory outputManagerFactory; - // NOTE: we use HashMap here to guarantee Serializability - // Mapping from view id to a view - private final HashMap> idToViewMap; - private final String transformFullName; - private final String transformId; - private final Coder inputCoder; - private final Coder> windowedValueCoder; - private final HashMap, Coder> outputCoders; - private final PCollection.IsBounded isBounded; - private final String bundleCheckTimerId; - private final String bundleStateId; - - // portable api related - private final boolean isPortable; - private final RunnerApi.ExecutableStagePayload stagePayload; - private final JobInfo jobInfo; - private final HashMap> idToTupleTagMap; - - private transient SamzaTimerInternalsFactory timerInternalsFactory; - private transient DoFnRunner fnRunner; - private transient PushbackSideInputDoFnRunner pushbackFnRunner; - private transient SideInputHandler sideInputHandler; - private transient DoFnInvoker doFnInvoker; - private transient SamzaPipelineOptions samzaPipelineOptions; - - // This is derivable from pushbackValues which is persisted to a store. - // TODO: eagerly initialize the hold in init - @edu.umd.cs.findbugs.annotations.SuppressWarnings( - justification = "No bug", - value = "SE_TRANSIENT_FIELD_NOT_RESTORED") - private transient Instant pushbackWatermarkHold; - - // TODO: add this to checkpointable state - private transient Instant inputWatermark; - private transient BundleManager bundleManager; - private transient Instant sideInputWatermark; - private transient List> pushbackValues; - private transient ExecutableStageContext stageContext; - private transient StageBundleFactory stageBundleFactory; - private transient boolean bundleDisabled; - - private final DoFnSchemaInformation doFnSchemaInformation; - private final Map> sideInputMapping; - private final Map stateIdToStoreMapping; - - public PortableDoFnOp( - TupleTag mainOutputTag, - DoFn doFn, - Coder keyCoder, - Coder inputCoder, - Coder> windowedValueCoder, - Map, Coder> outputCoders, - Collection> sideInputs, - List> sideOutputTags, - WindowingStrategy windowingStrategy, - Map> idToViewMap, - OutputManagerFactory outputManagerFactory, - String transformFullName, - String transformId, - PCollection.IsBounded isBounded, - boolean isPortable, - RunnerApi.ExecutableStagePayload stagePayload, - JobInfo jobInfo, - Map> idToTupleTagMap, - DoFnSchemaInformation doFnSchemaInformation, - Map> sideInputMapping, - Map stateIdToStoreMapping) { - this.mainOutputTag = mainOutputTag; - this.doFn = doFn; - this.sideInputs = sideInputs; - this.sideOutputTags = sideOutputTags; - this.inputCoder = inputCoder; - this.windowedValueCoder = windowedValueCoder; - this.outputCoders = new HashMap<>(outputCoders); - this.windowingStrategy = windowingStrategy; - this.idToViewMap = new HashMap<>(idToViewMap); - this.outputManagerFactory = outputManagerFactory; - this.transformFullName = transformFullName; - this.transformId = transformId; - this.keyCoder = keyCoder; - this.isBounded = isBounded; - this.isPortable = isPortable; - this.stagePayload = stagePayload; - this.jobInfo = jobInfo; - this.idToTupleTagMap = new HashMap<>(idToTupleTagMap); - this.bundleCheckTimerId = "_samza_bundle_check_" + transformId; - this.bundleStateId = "_samza_bundle_" + transformId; - this.doFnSchemaInformation = doFnSchemaInformation; - this.sideInputMapping = sideInputMapping; - this.stateIdToStoreMapping = stateIdToStoreMapping; - } - - @Override - @SuppressWarnings("unchecked") - public void open( - Config config, - Context context, - Scheduler> timerRegistry, - OpEmitter emitter) { - this.inputWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE; - this.sideInputWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE; - this.pushbackWatermarkHold = BoundedWindow.TIMESTAMP_MAX_VALUE; - - final DoFnSignature signature = DoFnSignatures.getSignature(doFn.getClass()); - final SamzaExecutionContext samzaExecutionContext = - (SamzaExecutionContext) context.getApplicationContainerContext(); - this.samzaPipelineOptions = samzaExecutionContext.getPipelineOptions(); - this.bundleDisabled = samzaPipelineOptions.getMaxBundleSize() <= 1; - - final String stateId = "pardo-" + transformId; - final SamzaStoreStateInternals.Factory nonKeyedStateInternalsFactory = - SamzaStoreStateInternals.createNonKeyedStateInternalsFactory( - stateId, context.getTaskContext(), samzaPipelineOptions); - final FutureCollector outputFutureCollector = createFutureCollector(); - - this.bundleManager = - new ClassicBundleManager<>( - createBundleProgressListener(), - outputFutureCollector, - samzaPipelineOptions.getMaxBundleSize(), - samzaPipelineOptions.getMaxBundleTimeMs(), - timerRegistry, - bundleCheckTimerId); - - this.timerInternalsFactory = - SamzaTimerInternalsFactory.createTimerInternalFactory( - keyCoder, - (Scheduler) timerRegistry, - getTimerStateId(signature), - nonKeyedStateInternalsFactory, - windowingStrategy, - isBounded, - samzaPipelineOptions); - - this.sideInputHandler = - new SideInputHandler(sideInputs, nonKeyedStateInternalsFactory.stateInternalsForKey(null)); - - if (isPortable) { - final ExecutableStage executableStage = ExecutableStage.fromPayload(stagePayload); - stageContext = SamzaExecutableStageContextFactory.getInstance().get(jobInfo); - stageBundleFactory = stageContext.getStageBundleFactory(executableStage); - this.fnRunner = - SamzaDoFnRunners.createPortable( - transformId, - DoFnUtils.toStepName(executableStage), - bundleStateId, - windowedValueCoder, - executableStage, - sideInputMapping, - sideInputHandler, - nonKeyedStateInternalsFactory, - timerInternalsFactory, - samzaPipelineOptions, - outputManagerFactory.create(emitter, outputFutureCollector), - stageBundleFactory, - samzaExecutionContext, - mainOutputTag, - idToTupleTagMap, - context, - transformFullName); - } else { - this.fnRunner = - SamzaDoFnRunners.create( - samzaPipelineOptions, - doFn, - windowingStrategy, - transformFullName, - stateId, - context, - mainOutputTag, - sideInputHandler, - timerInternalsFactory, - keyCoder, - outputManagerFactory.create(emitter, outputFutureCollector), - inputCoder, - sideOutputTags, - outputCoders, - doFnSchemaInformation, - (Map>) sideInputMapping, - stateIdToStoreMapping, - emitter, - outputFutureCollector); - } - - this.pushbackFnRunner = - SimplePushbackSideInputDoFnRunner.create(fnRunner, sideInputs, sideInputHandler); - this.pushbackValues = new ArrayList<>(); - - final Iterator invokerReg = - ServiceLoader.load(SamzaDoFnInvokerRegistrar.class).iterator(); - if (!invokerReg.hasNext()) { - // use the default invoker here - doFnInvoker = DoFnInvokers.tryInvokeSetupFor(doFn, samzaPipelineOptions); - } else { - doFnInvoker = - Iterators.getOnlyElement(invokerReg).invokerSetupFor(doFn, samzaPipelineOptions, context); - } - } - - FutureCollector createFutureCollector() { - return new FutureCollectorImpl<>(); - } - - private String getTimerStateId(DoFnSignature signature) { - final StringBuilder builder = new StringBuilder("timer"); - if (signature.usesTimers()) { - signature.timerDeclarations().keySet().forEach(builder::append); - } - return builder.toString(); - } - - @Override - public void processElement(WindowedValue inputElement, OpEmitter emitter) { - try { - bundleManager.tryStartBundle(); - final Iterable> rejectedValues = - pushbackFnRunner.processElementInReadyWindows(inputElement); - for (WindowedValue rejectedValue : rejectedValues) { - if (rejectedValue.getTimestamp().compareTo(pushbackWatermarkHold) < 0) { - pushbackWatermarkHold = rejectedValue.getTimestamp(); - } - pushbackValues.add(rejectedValue); - } - - bundleManager.tryFinishBundle(emitter); - } catch (Throwable t) { - LOG.error("Encountered error during process element", t); - bundleManager.signalFailure(t); - throw t; - } - } - - private void doProcessWatermark(Instant watermark, OpEmitter emitter) { - this.inputWatermark = watermark; - - if (sideInputWatermark.isEqual(BoundedWindow.TIMESTAMP_MAX_VALUE)) { - // this means we will never see any more side input - emitAllPushbackValues(); - } - - final Instant actualInputWatermark = - pushbackWatermarkHold.isBefore(inputWatermark) ? pushbackWatermarkHold : inputWatermark; - - timerInternalsFactory.setInputWatermark(actualInputWatermark); - - Collection> readyTimers = timerInternalsFactory.removeReadyTimers(); - if (!readyTimers.isEmpty()) { - pushbackFnRunner.startBundle(); - for (KeyedTimerData keyedTimerData : readyTimers) { - fireTimer(keyedTimerData); - } - pushbackFnRunner.finishBundle(); - } - - if (timerInternalsFactory.getOutputWatermark() == null - || timerInternalsFactory.getOutputWatermark().isBefore(actualInputWatermark)) { - timerInternalsFactory.setOutputWatermark(actualInputWatermark); - emitter.emitWatermark(timerInternalsFactory.getOutputWatermark()); - } - } - - @Override - public void processWatermark(Instant watermark, OpEmitter emitter) { - bundleManager.processWatermark(watermark, emitter); - } - - @Override - public void processSideInput( - String id, WindowedValue> elements, OpEmitter emitter) { - checkState( - bundleDisabled, "Side input not supported in bundling mode. Please disable bundling."); - @SuppressWarnings("unchecked") - final WindowedValue> retypedElements = (WindowedValue>) elements; - - final PCollectionView view = idToViewMap.get(id); - if (view == null) { - throw new IllegalArgumentException("No mapping of id " + id + " to view."); - } - - sideInputHandler.addSideInputValue(view, retypedElements); - - final List> previousPushbackValues = new ArrayList<>(pushbackValues); - pushbackWatermarkHold = BoundedWindow.TIMESTAMP_MAX_VALUE; - pushbackValues.clear(); - - for (final WindowedValue value : previousPushbackValues) { - processElement(value, emitter); - } - - // We may be able to advance the output watermark since we may have played some pushed back - // events. - processWatermark(this.inputWatermark, emitter); - } - - @Override - public void processSideInputWatermark(Instant watermark, OpEmitter emitter) { - checkState( - bundleDisabled, "Side input not supported in bundling mode. Please disable bundling."); - sideInputWatermark = watermark; - - if (sideInputWatermark.isEqual(BoundedWindow.TIMESTAMP_MAX_VALUE)) { - // this means we will never see any more side input - processWatermark(this.inputWatermark, emitter); - } - } - - @Override - @SuppressWarnings("unchecked") - public void processTimer(KeyedTimerData keyedTimerData, OpEmitter emitter) { - // this is internal timer in processing time to check whether a bundle should be closed - if (bundleCheckTimerId.equals(keyedTimerData.getTimerData().getTimerId())) { - bundleManager.processTimer(keyedTimerData, emitter); - return; - } - - pushbackFnRunner.startBundle(); - fireTimer(keyedTimerData); - pushbackFnRunner.finishBundle(); - - this.timerInternalsFactory.removeProcessingTimer((KeyedTimerData) keyedTimerData); - } - - @Override - public void close() { - doFnInvoker.invokeTeardown(); - try (AutoCloseable factory = stageBundleFactory; - AutoCloseable context = stageContext) { - // do nothing - } catch (Exception e) { - LOG.error("Failed to close stage bundle factory", e); - } - } - - private void fireTimer(KeyedTimerData keyedTimerData) { - final TimerInternals.TimerData timer = keyedTimerData.getTimerData(); - LOG.debug("Firing timer {}", timer); - - final StateNamespace namespace = timer.getNamespace(); - // NOTE: not sure why this is safe, but DoFnOperator makes this assumption - final BoundedWindow window = ((StateNamespaces.WindowNamespace) namespace).getWindow(); - - fnRunner.onTimer( - timer.getTimerId(), - timer.getTimerFamilyId(), - keyedTimerData.getKey(), - window, - timer.getTimestamp(), - timer.getOutputTimestamp(), - timer.getDomain(), - timer.causedByDrain()); - } - - // todo: should this go through bundle manager to start and finish the bundle? - private void emitAllPushbackValues() { - if (!pushbackValues.isEmpty()) { - pushbackFnRunner.startBundle(); - - final List> previousPushbackValues = new ArrayList<>(pushbackValues); - pushbackWatermarkHold = BoundedWindow.TIMESTAMP_MAX_VALUE; - pushbackValues.clear(); - - for (final WindowedValue value : previousPushbackValues) { - fnRunner.processElement(value); - } - - pushbackFnRunner.finishBundle(); - } - } - - private BundleManager.BundleProgressListener createBundleProgressListener() { - return new BundleManager.BundleProgressListener() { - @Override - public void onBundleStarted() { - pushbackFnRunner.startBundle(); - } - - @Override - public void onBundleFinished(OpEmitter emitter) { - pushbackFnRunner.finishBundle(); - } - - @Override - public void onWatermark(Instant watermark, OpEmitter emitter) { - doProcessWatermark(watermark, emitter); - } - }; - } -} From eda0099cb8502a9c566ee944ab813d790f3868a2 Mon Sep 17 00:00:00 2001 From: Reuven Lax Date: Thu, 28 May 2026 13:35:53 -0700 Subject: [PATCH 03/11] add Override --- .../main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java | 1 + 1 file changed, 1 insertion(+) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java index 100a392a9ed7..3e4675ab074a 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java @@ -2819,6 +2819,7 @@ public Instant fireTimestamp(DoFn doFn) { return currentTimer.getFireTimestamp(); } + @Override public @Nullable Object sideInput(String tagId) { PCollectionView view = sideInputMapping.get(tagId); if (view == null) { From 4a4348931ddcc5f3a84d618fb09d1574de8fd7c3 Mon Sep 17 00:00:00 2001 From: Reuven Lax Date: Sat, 30 May 2026 13:56:03 -0700 Subject: [PATCH 04/11] foo --- .../worker/CombineValuesFnFactory.java | 2 +- .../worker/SimpleDoFnRunnerFactory.java | 22 ++++++++++++++++++- .../dataflow/worker/UserParDoFnFactory.java | 2 +- 3 files changed, 23 insertions(+), 3 deletions(-) diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/CombineValuesFnFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/CombineValuesFnFactory.java index 50d518f78216..668ff7d297ba 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/CombineValuesFnFactory.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/CombineValuesFnFactory.java @@ -97,7 +97,7 @@ public ParDoFn create( operationContext, doFnInfo.getDoFnSchemaInformation(), doFnInfo.getSideInputMapping(), - SimpleDoFnRunnerFactory.INSTANCE); + SimpleDoFnRunnerFactory.INSTANCE_DONT_DELEGATE_SIDE_INPUTS); } private static DoFnInfo getDoFnInfo( diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleDoFnRunnerFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleDoFnRunnerFactory.java index 5286fc1aae90..7f7cdc8b9ab2 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleDoFnRunnerFactory.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleDoFnRunnerFactory.java @@ -24,6 +24,7 @@ import org.apache.beam.runners.core.SideInputReader; import org.apache.beam.sdk.coders.Coder; 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.DoFnSchemaInformation; import org.apache.beam.sdk.util.WindowedValueMultiReceiver; @@ -35,7 +36,13 @@ "rawtypes" // TODO(https://github.com/apache/beam/issues/20447) }) class SimpleDoFnRunnerFactory implements DoFnRunnerFactory { - public static final SimpleDoFnRunnerFactory INSTANCE = new SimpleDoFnRunnerFactory(); + private final boolean delegateStreamingSideInputs; + public static final SimpleDoFnRunnerFactory INSTANCE = new SimpleDoFnRunnerFactory(true); + public static final SimpleDoFnRunnerFactory INSTANCE_DONT_DELEGATE_SIDE_INPUTS = new SimpleDoFnRunnerFactory(false); + + public SimpleDoFnRunnerFactory(boolean delegateStreamingSideInputs) { + this.delegateStreamingSideInputs = delegateStreamingSideInputs; + } @Override public DoFnRunner createRunner( @@ -67,6 +74,19 @@ public DoFnRunner createRunner( windowingStrategy, doFnSchemaInformation, sideInputMapping); + if (delegateStreamingSideInputs) { + boolean hasStreamingSideInput = + options.as(StreamingOptions.class).isStreaming() && !sideInputReader.isEmpty(); + if (hasStreamingSideInput) { + return new StreamingSideInputDoFnRunner<>( + fnRunner, + new StreamingSideInputFetcher<>( + sideInputViews, + inputCoder, + windowingStrategy, + (StreamingModeExecutionContext.StreamingModeStepContext) userStepContext)); + } + } return fnRunner; } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UserParDoFnFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UserParDoFnFactory.java index a8d5975e45ea..ffeb2aa35460 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UserParDoFnFactory.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UserParDoFnFactory.java @@ -52,7 +52,7 @@ }) class UserParDoFnFactory implements ParDoFnFactory { static UserParDoFnFactory createDefault() { - return new UserParDoFnFactory(new UserDoFnExtractor(), SimpleDoFnRunnerFactory.INSTANCE); + return new UserParDoFnFactory(new UserDoFnExtractor(), SimpleDoFnRunnerFactory.INSTANCE_DONT_DELEGATE_SIDE_INPUTS); } interface DoFnExtractor { From 039f9a1b5b0ebe6ad607b75b32aa74124ca071b2 Mon Sep 17 00:00:00 2001 From: Reuven Lax Date: Wed, 3 Jun 2026 15:09:16 -0700 Subject: [PATCH 05/11] foo --- .../worker/CombineValuesFnFactory.java | 2 +- .../worker/SimpleDoFnRunnerFactory.java | 22 +- .../dataflow/worker/SimpleParDoFn.java | 506 +++--------------- .../dataflow/worker/SimpleParDoFnHelpers.java | 506 ++++++++++++++++++ .../worker/SplittableProcessFnFactory.java | 20 +- ...reamingKeyedWorkKitemSideInputParDoFn.java | 251 +++++++++ .../worker/StreamingSideInputFetcher.java | 2 +- .../worker/StreamingSideInputProcessor.java | 153 ++++-- .../dataflow/worker/UserParDoFnFactory.java | 56 +- .../dataflow/worker/SimpleParDoFnTest.java | 4 +- ...ingKeyedWorkKitemSideInputParDoFnTest.java | 494 +++++++++++++++++ .../worker/UserParDoFnFactoryTest.java | 10 +- 12 files changed, 1468 insertions(+), 558 deletions(-) create mode 100644 runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFnHelpers.java create mode 100644 runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkKitemSideInputParDoFn.java create mode 100644 runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkKitemSideInputParDoFnTest.java diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/CombineValuesFnFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/CombineValuesFnFactory.java index 668ff7d297ba..50d518f78216 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/CombineValuesFnFactory.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/CombineValuesFnFactory.java @@ -97,7 +97,7 @@ public ParDoFn create( operationContext, doFnInfo.getDoFnSchemaInformation(), doFnInfo.getSideInputMapping(), - SimpleDoFnRunnerFactory.INSTANCE_DONT_DELEGATE_SIDE_INPUTS); + SimpleDoFnRunnerFactory.INSTANCE); } private static DoFnInfo getDoFnInfo( diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleDoFnRunnerFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleDoFnRunnerFactory.java index 7f7cdc8b9ab2..5286fc1aae90 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleDoFnRunnerFactory.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleDoFnRunnerFactory.java @@ -24,7 +24,6 @@ import org.apache.beam.runners.core.SideInputReader; import org.apache.beam.sdk.coders.Coder; 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.DoFnSchemaInformation; import org.apache.beam.sdk.util.WindowedValueMultiReceiver; @@ -36,13 +35,7 @@ "rawtypes" // TODO(https://github.com/apache/beam/issues/20447) }) class SimpleDoFnRunnerFactory implements DoFnRunnerFactory { - private final boolean delegateStreamingSideInputs; - public static final SimpleDoFnRunnerFactory INSTANCE = new SimpleDoFnRunnerFactory(true); - public static final SimpleDoFnRunnerFactory INSTANCE_DONT_DELEGATE_SIDE_INPUTS = new SimpleDoFnRunnerFactory(false); - - public SimpleDoFnRunnerFactory(boolean delegateStreamingSideInputs) { - this.delegateStreamingSideInputs = delegateStreamingSideInputs; - } + public static final SimpleDoFnRunnerFactory INSTANCE = new SimpleDoFnRunnerFactory(); @Override public DoFnRunner createRunner( @@ -74,19 +67,6 @@ public DoFnRunner createRunner( windowingStrategy, doFnSchemaInformation, sideInputMapping); - if (delegateStreamingSideInputs) { - boolean hasStreamingSideInput = - options.as(StreamingOptions.class).isStreaming() && !sideInputReader.isEmpty(); - if (hasStreamingSideInput) { - return new StreamingSideInputDoFnRunner<>( - fnRunner, - new StreamingSideInputFetcher<>( - sideInputViews, - inputCoder, - windowingStrategy, - (StreamingModeExecutionContext.StreamingModeStepContext) userStepContext)); - } - } return fnRunner; } } 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 7203cbcae305..56f47dbcf067 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 @@ -17,56 +17,26 @@ */ package org.apache.beam.runners.dataflow.worker; -import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; -import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; - import java.io.Closeable; import java.util.Collection; import java.util.Collections; -import java.util.HashMap; import java.util.Iterator; -import java.util.List; import java.util.Map; -import org.apache.beam.runners.core.DoFnRunner; import org.apache.beam.runners.core.SideInputReader; -import org.apache.beam.runners.core.StateInternals; -import org.apache.beam.runners.core.StateNamespaces.WindowNamespace; -import org.apache.beam.runners.core.StateTag; -import org.apache.beam.runners.core.StateTags; -import org.apache.beam.runners.core.TimerInternals.TimerData; -import org.apache.beam.runners.dataflow.options.DataflowPipelineDebugOptions; -import org.apache.beam.runners.dataflow.worker.counters.Counter; -import org.apache.beam.runners.dataflow.worker.counters.CounterFactory; -import org.apache.beam.runners.dataflow.worker.counters.CounterName; -import org.apache.beam.runners.dataflow.worker.util.common.worker.ElementCounter; -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.Receiver; import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.options.PipelineOptions; -import org.apache.beam.sdk.options.StreamingOptions; -import org.apache.beam.sdk.state.StateSpec; -import org.apache.beam.sdk.state.TimeDomain; import org.apache.beam.sdk.transforms.DoFnSchemaInformation; -import org.apache.beam.sdk.transforms.reflect.DoFnSignature; -import org.apache.beam.sdk.transforms.reflect.DoFnSignature.StateDeclaration; -import org.apache.beam.sdk.transforms.reflect.DoFnSignatures; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; -import org.apache.beam.sdk.transforms.windowing.GlobalWindow; import org.apache.beam.sdk.util.DoFnInfo; -import org.apache.beam.sdk.util.WindowedValueMultiReceiver; import org.apache.beam.sdk.values.PCollectionView; 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.ImmutableList; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; import org.checkerframework.checker.nullness.qual.Nullable; -import org.joda.time.Duration; -import org.joda.time.Instant; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; /** * A base class providing simple set up, processing, and tear down for a wrapped {@link @@ -80,41 +50,7 @@ "nullness" // TODO(https://github.com/apache/beam/issues/20497) }) public class SimpleParDoFn implements ParDoFn { - - // TODO: Remove once Distributions has shipped. - @VisibleForTesting - static final String OUTPUTS_PER_ELEMENT_EXPERIMENT = "outputs_per_element_counter"; - - private static final String COUNTER_NAME = "per-element-output-count"; - - private static final Logger LOG = LoggerFactory.getLogger(SimpleParDoFn.class); - - protected final PipelineOptions options; - private final DoFnInstanceManager doFnInstanceManager; - - private final SideInputReader sideInputReader; - private final DataflowOperationContext operationContext; - private final TupleTag mainOutputTag; - private final Map, Integer> outputTupleTagsToReceiverIndices; - private final List> sideOutputTags; - private final DataflowExecutionContext.DataflowStepContext stepContext; - private final DataflowExecutionContext.DataflowStepContext userStepContext; - private final CounterFactory counterFactory; - private final DoFnRunnerFactory runnerFactory; - private final boolean hasStreamingSideInput; - private final OutputsPerElementTracker outputsPerElementTracker; - private final DoFnSchemaInformation doFnSchemaInformation; - private final Map> sideInputMapping; - - // Various DoFn helpers, null between bundles - private @Nullable DoFnRunner fnRunner; - @Nullable DoFnInfo fnInfo; - private Receiver @Nullable [] receivers; - - // This may additionally be null if it is not a real DoFn but an OldDoFn or - // GroupAlsoByWindowViaWindowSetDoFn - private @Nullable DoFnSignature fnSignature; - + private final SimpleParDoFnHelpers helpers; private @Nullable StreamingSideInputProcessor sideInputProcessor; /** Creates a {@link SimpleParDoFn} using basic information about the step being executed. */ @@ -129,220 +65,53 @@ public class SimpleParDoFn implements DoFnSchemaInformation doFnSchemaInformation, Map> sideInputMapping, DoFnRunnerFactory runnerFactory) { - this.options = options; - this.doFnInstanceManager = doFnInstanceManager; - - // We vend a freshly deserialized version for each run - this.sideInputReader = sideInputReader; - this.operationContext = operationContext; - checkArgument(!outputTupleTagsToReceiverIndices.isEmpty(), "expected at least one output"); - this.mainOutputTag = mainOutputTag; - this.outputTupleTagsToReceiverIndices = outputTupleTagsToReceiverIndices; - ImmutableList.Builder> sideOutputTagsBuilder = ImmutableList.builder(); - for (TupleTag tag : outputTupleTagsToReceiverIndices.keySet()) { - if (!mainOutputTag.equals(tag)) { - sideOutputTagsBuilder.add(tag); - } - } - this.sideOutputTags = sideOutputTagsBuilder.build(); - this.stepContext = stepContext; - - // StepContext provides a TimerInternals and StateInternals for use by the system - this class. - // For the user, we request a user-scoped StepContext to provide a user-scoped - // StateInternals and TimerInternals. - this.userStepContext = stepContext.namespacedToUser(); - - this.counterFactory = operationContext.counterFactory(); - this.runnerFactory = runnerFactory; - this.hasStreamingSideInput = - options.as(StreamingOptions.class).isStreaming() && !sideInputReader.isEmpty(); - this.outputsPerElementTracker = createOutputsPerElementTracker(); - this.doFnSchemaInformation = doFnSchemaInformation; - this.sideInputMapping = sideInputMapping; - } - - private OutputsPerElementTracker createOutputsPerElementTracker() { - // TODO: Remove once Distributions has shipped. - if (!hasExperiment(OUTPUTS_PER_ELEMENT_EXPERIMENT)) { - return NoopOutputsPerElementTracker.INSTANCE; - } - - // TODO: Remove log statement when functionality is enabled by default. - LOG.info("{} counter enabled.", COUNTER_NAME); - - return new OutputsPerElementTrackerImpl(); - } - - private boolean hasExperiment(String experiment) { - List experiments = options.as(DataflowPipelineDebugOptions.class).getExperiments(); - return experiments != null && experiments.contains(experiment); - } - - /** Simple state tracker to calculate PerElementOutputCount counter. */ - private interface OutputsPerElementTracker { - - void onOutput(); - - void onProcessElement(); - - void onProcessElementSuccess(); - } - - private class OutputsPerElementTrackerImpl implements OutputsPerElementTracker { - - private long outputsPerElement; - private final Counter counter; - - public OutputsPerElementTrackerImpl() { - this.counter = - counterFactory.distribution( - CounterName.named(COUNTER_NAME).withOriginalName(stepContext.getNameContext())); - } - - @Override - public void onProcessElement() { - reset(); - } - - @Override - public void onOutput() { - outputsPerElement++; - } - - @Override - public void onProcessElementSuccess() { - counter.addValue(outputsPerElement); - reset(); - } - - private void reset() { - outputsPerElement = 0L; - } - } - - /** No-op {@link OutputsPerElementTracker} implementation used when the counter is disabled. */ - private static class NoopOutputsPerElementTracker implements OutputsPerElementTracker { - - private NoopOutputsPerElementTracker() {} - - public static final OutputsPerElementTracker INSTANCE = new NoopOutputsPerElementTracker(); - - @Override - public void onOutput() {} - - @Override - public void onProcessElement() {} - - @Override - public void onProcessElementSuccess() {} + this.helpers = + new SimpleParDoFnHelpers<>( + options, + doFnInstanceManager, + sideInputReader, + mainOutputTag, + outputTupleTagsToReceiverIndices, + stepContext, + operationContext, + doFnSchemaInformation, + sideInputMapping, + runnerFactory); } @Override public void startBundle(Receiver... receivers) throws Exception { - checkArgument( - receivers.length == outputTupleTagsToReceiverIndices.size(), - "unexpected number of receivers for DoFn"); - - this.receivers = receivers; - if (hasStreamingSideInput) { + helpers.startBundle(receivers); + if (helpers.hasStreamingSideInput) { // There is non-trivial setup that needs to be performed for watermark propagation // even on empty bundles. - reallyStartBundle(); + helpers.reallyStartBundle(); + onStartKey(); } } - private void reallyStartBundle() throws Exception { - checkState(fnRunner == null, "bundle already started (or not properly finished)"); - - WindowedValueMultiReceiver outputManager = - new WindowedValueMultiReceiver() { - final Map, OutputReceiver> undeclaredOutputs = new HashMap<>(); - - private @Nullable Receiver getReceiverOrNull(TupleTag tag) { - Integer receiverIndex = outputTupleTagsToReceiverIndices.get(tag); - if (receiverIndex != null) { - return receivers[receiverIndex]; - } else { - return undeclaredOutputs.get(tag); - } - } - - @Override - public void output(TupleTag tag, WindowedValue output) { - outputsPerElementTracker.onOutput(); - Receiver receiver = getReceiverOrNull(tag); - if (receiver == null) { - // A new undeclared output. - // TODO: plumb through the operationName, so that we can - // name implicit outputs after it. - String outputName = "implicit-" + tag.getId(); - // TODO: plumb through the counter prefix, so we can - // make it available to the OutputReceiver class in case - // it wants to use it in naming output counterFactory. (It - // doesn't today.) - OutputReceiver undeclaredReceiver = new OutputReceiver(); - - ElementCounter outputCounter = - new DataflowOutputCounter( - outputName, counterFactory, stepContext.getNameContext()); - undeclaredReceiver.addOutputCounter(outputCounter); - undeclaredOutputs.put(tag, undeclaredReceiver); - receiver = undeclaredReceiver; - } - - try { - receiver.process(output); - } catch (RuntimeException | Error e) { - // Rethrow unchecked exceptions as-is to avoid excessive nesting - // via a chain of DoFn's. - throw e; - } catch (Exception e) { - // This should never happen in practice with DoFn's, but can happen - // with other Receivers. - throw new RuntimeException(e); - } - } - }; - fnInfo = (DoFnInfo) doFnInstanceManager.get(); - fnSignature = DoFnSignatures.getSignature(fnInfo.getDoFn().getClass()); - - fnRunner = - runnerFactory.createRunner( - fnInfo.getDoFn(), - options, - mainOutputTag, - sideOutputTags, - fnInfo.getSideInputViews(), - sideInputReader, - fnInfo.getInputCoder(), - fnInfo.getOutputCoders(), - fnInfo.getWindowingStrategy(), - stepContext, - userStepContext, - outputManager, - doFnSchemaInformation, - sideInputMapping); - if (hasStreamingSideInput) { + protected void onStartKey() { + // TODO(relax): This assumes single-key bundles, which will change! Refactor this to not make + // this assumption. + if (helpers.hasStreamingSideInput) { sideInputProcessor = new StreamingSideInputProcessor<>( new StreamingSideInputFetcher( - fnInfo.getSideInputViews(), - fnInfo.getInputCoder(), - (WindowingStrategy) fnInfo.getWindowingStrategy(), - (StreamingModeExecutionContext.StreamingModeStepContext) userStepContext)); - } + helpers.fnInfo.getSideInputViews(), + helpers.fnInfo.getInputCoder(), + (WindowingStrategy) helpers.fnInfo.getWindowingStrategy(), + (StreamingModeExecutionContext.StreamingModeStepContext) + helpers.userStepContext)); - fnRunner.startBundle(); - if (sideInputProcessor != null) { - boolean hasState = fnSignature != null && !fnSignature.stateDeclarations().isEmpty(); + boolean hasState = helpers.hasState(); Iterator> unblockedElements = sideInputProcessor.tryUnblockElements(); for (Iterator> it = unblockedElements; it.hasNext(); ) { WindowedValue unblockedElement = it.next(); - fnRunner.processElement(unblockedElement); + helpers.fnRunner.processElement(unblockedElement); if (hasState) { // These elements are now processed. Register cleanup timers for all the unblocked // windows. - registerStateCleanup( + helpers.registerStateCleanup( (WindowingStrategy) getDoFnInfo().getWindowingStrategy(), (Collection) unblockedElement.getWindows()); } @@ -353,18 +122,24 @@ public void output(TupleTag tag, WindowedValue output) { @Override @SuppressWarnings("unchecked") public void processElement(Object untypedElem) throws Exception { - if (fnRunner == null) { + if (helpers.fnRunner == null) { // If we need to run reallyStartBundle in here, we need to make sure to switch the state // sampler into the start state. - try (Closeable start = operationContext.enterStart()) { - reallyStartBundle(); + try (Closeable start = helpers.operationContext.enterStart()) { + helpers.reallyStartBundle(); + onStartKey(); } } + helpers.outputsPerElementTracker.onProcessElement(); WindowedValue elem = (WindowedValue) untypedElem; + onProcessWindowedValue(elem); + + helpers.outputsPerElementTracker.onProcessElementSuccess(); + } - boolean hasState = fnSignature != null && !fnSignature.stateDeclarations().isEmpty(); - outputsPerElementTracker.onProcessElement(); + protected void onProcessWindowedValue(WindowedValue elem) { + boolean hasState = helpers.hasState(); Collection windowsProcessed; if (sideInputProcessor != null) { @@ -373,7 +148,7 @@ public void processElement(Object untypedElem) throws Exception { sideInputProcessor.handleProcessElement(elem); it.hasNext(); ) { WindowedValue toProcess = it.next(); - fnRunner.processElement(toProcess); + helpers.fnRunner.processElement(toProcess); if (hasState) { windowsProcessed.addAll((Collection) toProcess.getWindows()); // If the element was blocked, don't register a cleanup timer. The timer will be @@ -382,14 +157,13 @@ public void processElement(Object untypedElem) throws Exception { } } } else { - fnRunner.processElement(elem); + helpers.fnRunner.processElement(elem); windowsProcessed = (Collection) elem.getWindows(); } if (hasState) { - registerStateCleanup( + helpers.registerStateCleanup( (WindowingStrategy) getDoFnInfo().getWindowingStrategy(), windowsProcessed); } - outputsPerElementTracker.onProcessElementSuccess(); } @Override @@ -402,193 +176,33 @@ public void processTimers() throws Exception { // exist without actually decoding them. Coder windowCoder = (Coder) - (fnInfo != null ? fnInfo : doFnInstanceManager.peek()) + (helpers.fnInfo != null ? helpers.fnInfo : helpers.doFnInstanceManager.peek()) .getWindowingStrategy() .getWindowFn() .windowCoder(); - processTimers(TimerType.USER, userStepContext, windowCoder); - processTimers(TimerType.SYSTEM, stepContext, windowCoder); - } - - private void processUserTimer(TimerData timer) throws Exception { - if (fnSignature.timerDeclarations().containsKey(timer.getTimerId()) - || fnSignature.timerFamilyDeclarations().containsKey(timer.getTimerFamilyId())) { - BoundedWindow window = ((WindowNamespace) timer.getNamespace()).getWindow(); - if (sideInputProcessor != null) { - sideInputProcessor.handleProcessTimer(timer); - } - fnRunner.onTimer( - timer.getTimerId(), - timer.getTimerFamilyId(), - this.stepContext.stateInternals().getKey(), - window, - timer.getTimestamp(), - timer.getOutputTimestamp(), - timer.getDomain(), - timer.causedByDrain()); - } - } - - private void processSystemTimer(TimerData timer) throws Exception { - // Timer owned by this class, for cleaning up state in expired windows - if (timer.getTimerId().equals(CLEANUP_TIMER_ID)) { - checkState( - timer.getDomain().equals(TimeDomain.EVENT_TIME), - "%s received cleanup timer with domain not EVENT_TIME: %s", - this, - timer); - - checkState( - timer.getNamespace() instanceof WindowNamespace, - "%s received cleanup timer not for a %s: %s", - this, - WindowNamespace.class.getSimpleName(), - timer); - - if (sideInputProcessor != null) { - // We must call this to ensure the side-input is cached for onWindowExpiration. Since we - // don't set cleanup - // timers until we actually call processElement, the window must be unblocked here. - sideInputProcessor.handleProcessTimer(timer); - } - - BoundedWindow window = ((WindowNamespace) timer.getNamespace()).getWindow(); - Instant targetTime = earliestAllowableCleanupTime(window, fnInfo.getWindowingStrategy()); - - checkState( - !targetTime.isAfter(timer.getTimestamp()), - "%s received state cleanup timer for window %s " - + " that is before the appropriate cleanup time %s", - this, - window, - targetTime); - - fnRunner.onWindowExpiration( - window, timer.getOutputTimestamp(), this.stepContext.stateInternals().getKey()); - - // This is for a timer for a window that is expired, so clean it up. - for (StateDeclaration stateDecl : fnSignature.stateDeclarations().values()) { - StateTag tag; - try { - tag = - StateTags.tagForSpec( - stateDecl.id(), (StateSpec) stateDecl.field().get(fnInfo.getDoFn())); - } catch (IllegalAccessException e) { - throw new RuntimeException( - String.format( - "Error accessing %s for %s", - StateSpec.class.getName(), fnInfo.getDoFn().getClass().getName()), - e); - } - - StateInternals stateInternals = userStepContext.stateInternals(); - org.apache.beam.sdk.state.State state = stateInternals.state(timer.getNamespace(), tag); - state.clear(); - } - } + helpers.processTimers( + SimpleParDoFnHelpers.TimerType.USER, + helpers.userStepContext, + windowCoder, + this::onStartKey, + sideInputProcessor); + helpers.processTimers( + SimpleParDoFnHelpers.TimerType.SYSTEM, + helpers.stepContext, + windowCoder, + this::onStartKey, + sideInputProcessor); } @Override public void finishBundle() throws Exception { - if (fnRunner != null) { - fnRunner.finishBundle(); - if (sideInputProcessor != null) { - sideInputProcessor.handleFinishBundle(); - } - doFnInstanceManager.complete(fnInfo); - fnRunner = null; - fnInfo = null; - fnSignature = null; - sideInputProcessor = null; - } + helpers.finishBundle(sideInputProcessor); + this.sideInputProcessor = null; } @Override public void abort() throws Exception { - doFnInstanceManager.abort(fnInfo); - fnRunner = null; - fnInfo = null; - } - - @VisibleForTesting static final String CLEANUP_TIMER_ID = "cleanup-timer"; - - private enum TimerType { - USER { - @Override - public void processTimer(SimpleParDoFn doFn, TimerData timer) throws Exception { - doFn.processUserTimer(timer); - } - }, - SYSTEM { - @Override - public void processTimer(SimpleParDoFn doFn, TimerData timer) throws Exception { - doFn.processSystemTimer(timer); - } - }; - - public abstract void processTimer(SimpleParDoFn doFn, TimerData timer) throws Exception; - }; - - private void processTimers( - TimerType mode, - DataflowExecutionContext.DataflowStepContext context, - Coder windowCoder) - throws Exception { - TimerData timer = context.getNextFiredTimer(windowCoder); - - if (timer != null && fnRunner == null) { - // If we need to run reallyStartBundle in here, we need to make sure to switch the state - // sampler into the start state. - try (Closeable start = operationContext.enterStart()) { - reallyStartBundle(); - } - } - - while (timer != null) { - mode.processTimer(this, timer); - timer = context.getNextFiredTimer(windowCoder); - } - } - - private void registerStateCleanup( - WindowingStrategy windowingStrategy, Collection windowsToCleanup) { - Coder windowCoder = windowingStrategy.getWindowFn().windowCoder(); - - for (W window : windowsToCleanup) { - // The stepContext is the thing that know if it is batch or streaming, hence - // whether state needs to be cleaned up or will simply be discarded so the - // timer can be ignored. - Instant cleanupTime = earliestAllowableCleanupTime(window, windowingStrategy); - // Set a cleanup timer for state at the end of the window to trigger onWindowExpiration and - // garbage collect state. We avoid doing this for the global window if there is no window - // expiration set as the state will be up when the pipeline terminates. Setting the timer - // leads to a unbounded growth of timers for pipelines with many unique keys in the global - // window. - if (cleanupTime.isBefore(GlobalWindow.INSTANCE.maxTimestamp()) - || fnSignature.onWindowExpiration() != null) { - // If the DoFn has OnWindowExpiration, then set the watermark hold so that the watermark - // does - // not advance until OnWindowExpiration completes. - Instant cleanupOutputTimestamp = - fnSignature.onWindowExpiration() == null - ? cleanupTime - : cleanupTime.minus(Duration.millis(1L)); - stepContext.setStateCleanupTimer( - CLEANUP_TIMER_ID, window, windowCoder, cleanupTime, cleanupOutputTimestamp); - } - } - } - - private Instant earliestAllowableCleanupTime( - BoundedWindow window, WindowingStrategy windowingStrategy) { - Instant cleanupTime = - window - .maxTimestamp() - .plus(windowingStrategy.getAllowedLateness()) - .plus(Duration.millis(1L)); - return cleanupTime.isAfter(BoundedWindow.TIMESTAMP_MAX_VALUE) - ? BoundedWindow.TIMESTAMP_MAX_VALUE - : cleanupTime; + helpers.abort(); } /** @@ -600,6 +214,6 @@ private Instant earliestAllowableCleanupTime( @VisibleForTesting @Nullable DoFnInfo getDoFnInfo() { - return fnInfo; + return helpers.fnInfo; } } 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 new file mode 100644 index 000000000000..6a969c97b06a --- /dev/null +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFnHelpers.java @@ -0,0 +1,506 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.dataflow.worker; + +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; + +import java.io.Closeable; +import java.util.Collection; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.apache.beam.runners.core.DoFnRunner; +import org.apache.beam.runners.core.SideInputReader; +import org.apache.beam.runners.core.StateInternals; +import org.apache.beam.runners.core.StateNamespaces; +import org.apache.beam.runners.core.StateTag; +import org.apache.beam.runners.core.StateTags; +import org.apache.beam.runners.core.TimerInternals; +import org.apache.beam.runners.dataflow.options.DataflowPipelineDebugOptions; +import org.apache.beam.runners.dataflow.worker.counters.Counter; +import org.apache.beam.runners.dataflow.worker.counters.CounterFactory; +import org.apache.beam.runners.dataflow.worker.counters.CounterName; +import org.apache.beam.runners.dataflow.worker.util.common.worker.ElementCounter; +import org.apache.beam.runners.dataflow.worker.util.common.worker.OutputReceiver; +import org.apache.beam.runners.dataflow.worker.util.common.worker.Receiver; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.StreamingOptions; +import org.apache.beam.sdk.state.StateSpec; +import org.apache.beam.sdk.state.TimeDomain; +import org.apache.beam.sdk.transforms.DoFnSchemaInformation; +import org.apache.beam.sdk.transforms.reflect.DoFnSignature; +import org.apache.beam.sdk.transforms.reflect.DoFnSignatures; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; +import org.apache.beam.sdk.util.DoFnInfo; +import org.apache.beam.sdk.util.WindowedValueMultiReceiver; +import org.apache.beam.sdk.values.PCollectionView; +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.ImmutableList; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Duration; +import org.joda.time.Instant; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +@SuppressWarnings({ + "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) + "nullness" // TODO(https://github.com/apache/beam/issues/20497) +}) +public class SimpleParDoFnHelpers { + private static final Logger LOG = LoggerFactory.getLogger(SimpleParDoFnHelpers.class); + + // TODO: Remove once Distributions has shipped. + @VisibleForTesting + static final String OUTPUTS_PER_ELEMENT_EXPERIMENT = "outputs_per_element_counter"; + + private static final String COUNTER_NAME = "per-element-output-count"; + + final PipelineOptions options; + final DoFnInstanceManager doFnInstanceManager; + + private final SideInputReader sideInputReader; + final DataflowOperationContext operationContext; + private final TupleTag mainOutputTag; + private final Map, Integer> outputTupleTagsToReceiverIndices; + private final List> sideOutputTags; + final DataflowExecutionContext.DataflowStepContext stepContext; + final DataflowExecutionContext.DataflowStepContext userStepContext; + private final CounterFactory counterFactory; + private final DoFnRunnerFactory runnerFactory; + final boolean hasStreamingSideInput; + final OutputsPerElementTracker outputsPerElementTracker; + private final DoFnSchemaInformation doFnSchemaInformation; + private final Map> sideInputMapping; + + // Various DoFn helpers, null between bundles + @Nullable DoFnRunner fnRunner; + @Nullable DoFnInfo fnInfo; + private Receiver @Nullable [] receivers; + + // This may additionally be null if it is not a real DoFn but an OldDoFn or + // GroupAlsoByWindowViaWindowSetDoFn + protected @Nullable DoFnSignature fnSignature; + + SimpleParDoFnHelpers( + PipelineOptions options, + DoFnInstanceManager doFnInstanceManager, + SideInputReader sideInputReader, + TupleTag mainOutputTag, + Map, Integer> outputTupleTagsToReceiverIndices, + DataflowExecutionContext.DataflowStepContext stepContext, + DataflowOperationContext operationContext, + DoFnSchemaInformation doFnSchemaInformation, + Map> sideInputMapping, + DoFnRunnerFactory runnerFactory) { + this.options = options; + this.doFnInstanceManager = doFnInstanceManager; + + // We vend a freshly deserialized version for each run + this.sideInputReader = sideInputReader; + this.operationContext = operationContext; + checkArgument(!outputTupleTagsToReceiverIndices.isEmpty(), "expected at least one output"); + this.mainOutputTag = mainOutputTag; + this.outputTupleTagsToReceiverIndices = outputTupleTagsToReceiverIndices; + ImmutableList.Builder> sideOutputTagsBuilder = ImmutableList.builder(); + for (TupleTag tag : outputTupleTagsToReceiverIndices.keySet()) { + if (!mainOutputTag.equals(tag)) { + sideOutputTagsBuilder.add(tag); + } + } + this.sideOutputTags = sideOutputTagsBuilder.build(); + this.stepContext = stepContext; + + // StepContext provides a TimerInternals and StateInternals for use by the system - this class. + // For the user, we request a user-scoped StepContext to provide a user-scoped + // StateInternals and TimerInternals. + this.userStepContext = stepContext.namespacedToUser(); + + this.counterFactory = operationContext.counterFactory(); + this.runnerFactory = runnerFactory; + this.hasStreamingSideInput = + options.as(StreamingOptions.class).isStreaming() && !sideInputReader.isEmpty(); + this.outputsPerElementTracker = createOutputsPerElementTracker(); + this.doFnSchemaInformation = doFnSchemaInformation; + this.sideInputMapping = sideInputMapping; + } + + boolean hasState() { + return fnSignature != null && !fnSignature.stateDeclarations().isEmpty(); + } + + void startBundle(Receiver... receivers) throws Exception { + checkArgument( + receivers.length == outputTupleTagsToReceiverIndices.size(), + "unexpected number of receivers for DoFn"); + + this.receivers = receivers; + } + + void reallyStartBundle() throws Exception { + checkState(fnRunner == null, "bundle already started (or not properly finished)"); + + WindowedValueMultiReceiver outputManager = + new WindowedValueMultiReceiver() { + final Map, OutputReceiver> undeclaredOutputs = new HashMap<>(); + + private @Nullable Receiver getReceiverOrNull(TupleTag tag) { + Integer receiverIndex = outputTupleTagsToReceiverIndices.get(tag); + if (receiverIndex != null) { + return receivers[receiverIndex]; + } else { + return undeclaredOutputs.get(tag); + } + } + + @Override + public void output(TupleTag tag, WindowedValue output) { + outputsPerElementTracker.onOutput(); + Receiver receiver = getReceiverOrNull(tag); + if (receiver == null) { + // A new undeclared output. + // TODO: plumb through the operationName, so that we can + // name implicit outputs after it. + String outputName = "implicit-" + tag.getId(); + // TODO: plumb through the counter prefix, so we can + // make it available to the OutputReceiver class in case + // it wants to use it in naming output counterFactory. (It + // doesn't today.) + OutputReceiver undeclaredReceiver = new OutputReceiver(); + + ElementCounter outputCounter = + new DataflowOutputCounter( + outputName, counterFactory, stepContext.getNameContext()); + undeclaredReceiver.addOutputCounter(outputCounter); + undeclaredOutputs.put(tag, undeclaredReceiver); + receiver = undeclaredReceiver; + } + + try { + receiver.process(output); + } catch (RuntimeException | Error e) { + // Rethrow unchecked exceptions as-is to avoid excessive nesting + // via a chain of DoFn's. + throw e; + } catch (Exception e) { + // This should never happen in practice with DoFn's, but can happen + // with other Receivers. + throw new RuntimeException(e); + } + } + }; + fnInfo = (DoFnInfo) doFnInstanceManager.get(); + fnSignature = DoFnSignatures.getSignature(fnInfo.getDoFn().getClass()); + + fnRunner = + runnerFactory.createRunner( + fnInfo.getDoFn(), + options, + mainOutputTag, + sideOutputTags, + fnInfo.getSideInputViews(), + sideInputReader, + fnInfo.getInputCoder(), + fnInfo.getOutputCoders(), + fnInfo.getWindowingStrategy(), + stepContext, + userStepContext, + outputManager, + doFnSchemaInformation, + sideInputMapping); + fnRunner.startBundle(); + } + + void finishBundle(StreamingSideInputProcessor sideInputProcessor) throws Exception { + if (fnRunner != null) { + fnRunner.finishBundle(); + if (sideInputProcessor != null) { + sideInputProcessor.handleFinishBundle(); + } + doFnInstanceManager.complete(fnInfo); + fnRunner = null; + fnInfo = null; + fnSignature = null; + sideInputProcessor = null; + } + } + + void abort() throws Exception { + doFnInstanceManager.abort(fnInfo); + fnRunner = null; + fnInfo = null; + } + + @VisibleForTesting static final String CLEANUP_TIMER_ID = "cleanup-timer"; + + enum TimerType { + USER { + @Override + public void processTimer( + SimpleParDoFnHelpers doFn, + TimerInternals.TimerData timer, + StreamingSideInputProcessor sideInputProcessor) + throws Exception { + doFn.processUserTimer(timer, sideInputProcessor); + } + }, + SYSTEM { + @Override + public void processTimer( + SimpleParDoFnHelpers doFn, + TimerInternals.TimerData timer, + StreamingSideInputProcessor sideInputProcessor) + throws Exception { + doFn.processSystemTimer(timer, sideInputProcessor); + } + }; + + public abstract void processTimer( + SimpleParDoFnHelpers doFn, + TimerInternals.TimerData timer, + StreamingSideInputProcessor sideInputProcessor) + throws Exception; + }; + + void processTimers( + TimerType mode, + DataflowExecutionContext.DataflowStepContext context, + Coder windowCoder, + Runnable startKey, + StreamingSideInputProcessor sideInputProcessor) + throws Exception { + TimerInternals.TimerData timer = context.getNextFiredTimer(windowCoder); + + if (timer != null && fnRunner == null) { + // If we need to run reallyStartBundle in here, we need to make sure to switch the state + // sampler into the start state. + try (Closeable start = operationContext.enterStart()) { + reallyStartBundle(); + startKey.run(); + } + } + + while (timer != null) { + mode.processTimer(this, timer, sideInputProcessor); + timer = context.getNextFiredTimer(windowCoder); + } + } + + protected void processUserTimer( + TimerInternals.TimerData timer, StreamingSideInputProcessor sideInputProcessor) + throws Exception { + if (fnSignature.timerDeclarations().containsKey(timer.getTimerId()) + || fnSignature.timerFamilyDeclarations().containsKey(timer.getTimerFamilyId())) { + BoundedWindow window = ((StateNamespaces.WindowNamespace) timer.getNamespace()).getWindow(); + if (sideInputProcessor != null) { + sideInputProcessor.handleProcessTimer(timer); + } + fnRunner.onTimer( + timer.getTimerId(), + timer.getTimerFamilyId(), + this.stepContext.stateInternals().getKey(), + window, + timer.getTimestamp(), + timer.getOutputTimestamp(), + timer.getDomain(), + timer.causedByDrain()); + } + } + + private void processSystemTimer( + TimerInternals.TimerData timer, StreamingSideInputProcessor sideInputProcessor) + throws Exception { + // Timer owned by this class, for cleaning up state in expired windows + if (timer.getTimerId().equals(CLEANUP_TIMER_ID)) { + checkState( + timer.getDomain().equals(TimeDomain.EVENT_TIME), + "%s received cleanup timer with domain not EVENT_TIME: %s", + this, + timer); + + checkState( + timer.getNamespace() instanceof StateNamespaces.WindowNamespace, + "%s received cleanup timer not for a %s: %s", + this, + StateNamespaces.WindowNamespace.class.getSimpleName(), + timer); + + if (sideInputProcessor != null) { + // We must call this to ensure the side-input is cached for onWindowExpiration. Since we + // don't set cleanup + // timers until we actually call processElement, the window must be unblocked here. + sideInputProcessor.handleProcessTimer(timer); + } + + BoundedWindow window = ((StateNamespaces.WindowNamespace) timer.getNamespace()).getWindow(); + Instant targetTime = earliestAllowableCleanupTime(window, fnInfo.getWindowingStrategy()); + + checkState( + !targetTime.isAfter(timer.getTimestamp()), + "%s received state cleanup timer for window %s " + + " that is before the appropriate cleanup time %s", + this, + window, + targetTime); + + fnRunner.onWindowExpiration( + window, timer.getOutputTimestamp(), this.stepContext.stateInternals().getKey()); + + // This is for a timer for a window that is expired, so clean it up. + for (DoFnSignature.StateDeclaration stateDecl : fnSignature.stateDeclarations().values()) { + StateTag tag; + try { + tag = + StateTags.tagForSpec( + stateDecl.id(), (StateSpec) stateDecl.field().get(fnInfo.getDoFn())); + } catch (IllegalAccessException e) { + throw new RuntimeException( + String.format( + "Error accessing %s for %s", + StateSpec.class.getName(), fnInfo.getDoFn().getClass().getName()), + e); + } + + StateInternals stateInternals = userStepContext.stateInternals(); + org.apache.beam.sdk.state.State state = stateInternals.state(timer.getNamespace(), tag); + state.clear(); + } + } + } + + private OutputsPerElementTracker createOutputsPerElementTracker() { + // TODO: Remove once Distributions has shipped. + if (!hasExperiment(OUTPUTS_PER_ELEMENT_EXPERIMENT)) { + return NoopOutputsPerElementTracker.INSTANCE; + } + + // TODO: Remove log statement when functionality is enabled by default. + LOG.info("{} counter enabled.", COUNTER_NAME); + + return new OutputsPerElementTrackerImpl(); + } + + private boolean hasExperiment(String experiment) { + List experiments = options.as(DataflowPipelineDebugOptions.class).getExperiments(); + return experiments != null && experiments.contains(experiment); + } + + /** Simple state tracker to calculate PerElementOutputCount counter. */ + interface OutputsPerElementTracker { + + void onOutput(); + + void onProcessElement(); + + void onProcessElementSuccess(); + } + + private class OutputsPerElementTrackerImpl implements OutputsPerElementTracker { + + private long outputsPerElement; + private final Counter counter; + + public OutputsPerElementTrackerImpl() { + this.counter = + counterFactory.distribution( + CounterName.named(COUNTER_NAME).withOriginalName(stepContext.getNameContext())); + } + + @Override + public void onProcessElement() { + reset(); + } + + @Override + public void onOutput() { + outputsPerElement++; + } + + @Override + public void onProcessElementSuccess() { + counter.addValue(outputsPerElement); + reset(); + } + + private void reset() { + outputsPerElement = 0L; + } + } + + /** No-op {@link OutputsPerElementTracker} implementation used when the counter is disabled. */ + private static class NoopOutputsPerElementTracker implements OutputsPerElementTracker { + + private NoopOutputsPerElementTracker() {} + + public static final OutputsPerElementTracker INSTANCE = new NoopOutputsPerElementTracker(); + + @Override + public void onOutput() {} + + @Override + public void onProcessElement() {} + + @Override + public void onProcessElementSuccess() {} + } + + Instant earliestAllowableCleanupTime(BoundedWindow window, WindowingStrategy windowingStrategy) { + Instant cleanupTime = + window + .maxTimestamp() + .plus(windowingStrategy.getAllowedLateness()) + .plus(Duration.millis(1L)); + return cleanupTime.isAfter(BoundedWindow.TIMESTAMP_MAX_VALUE) + ? BoundedWindow.TIMESTAMP_MAX_VALUE + : cleanupTime; + } + + protected void registerStateCleanup( + WindowingStrategy windowingStrategy, Collection windowsToCleanup) { + Coder windowCoder = windowingStrategy.getWindowFn().windowCoder(); + + for (W window : windowsToCleanup) { + // The stepContext is the thing that know if it is batch or streaming, hence + // whether state needs to be cleaned up or will simply be discarded so the + // timer can be ignored. + Instant cleanupTime = earliestAllowableCleanupTime(window, windowingStrategy); + // Set a cleanup timer for state at the end of the window to trigger onWindowExpiration and + // garbage collect state. We avoid doing this for the global window if there is no window + // expiration set as the state will be up when the pipeline terminates. Setting the timer + // leads to a unbounded growth of timers for pipelines with many unique keys in the global + // window. + if (cleanupTime.isBefore(GlobalWindow.INSTANCE.maxTimestamp()) + || fnSignature.onWindowExpiration() != null) { + // If the DoFn has OnWindowExpiration, then set the watermark hold so that the watermark + // does + // not advance until OnWindowExpiration completes. + Instant cleanupOutputTimestamp = + fnSignature.onWindowExpiration() == null + ? cleanupTime + : cleanupTime.minus(Duration.millis(1L)); + stepContext.setStateCleanupTimer( + CLEANUP_TIMER_ID, window, windowCoder, cleanupTime, cleanupOutputTimestamp); + } + } + } +} 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 3ad443ee2a2b..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 @@ -44,7 +44,6 @@ import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.coders.KvCoder; 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.DoFnSchemaInformation; import org.apache.beam.sdk.util.DoFnInfo; @@ -65,7 +64,8 @@ }) class SplittableProcessFnFactory { static final ParDoFnFactory createDefault() { - return new UserParDoFnFactory(new ProcessFnExtractor(), new SplittableDoFnRunnerFactory()); + return new UserParDoFnFactory( + new ProcessFnExtractor(), new SplittableDoFnRunnerFactory(), true); } private static class ProcessFnExtractor implements UserParDoFnFactory.DoFnExtractor { @@ -174,22 +174,6 @@ public DoFnRunner>, OutputT> crea sideInputMapping); DoFnRunner>, OutputT> fnRunner = new DataflowProcessFnRunner<>(simpleRunner); - boolean hasStreamingSideInput = - options.as(StreamingOptions.class).isStreaming() && !sideInputReader.isEmpty(); - KeyedWorkItemCoder> kwiCoder = - (KeyedWorkItemCoder>) inputCoder; - if (hasStreamingSideInput) { - fnRunner = - new StreamingKeyedWorkItemSideInputDoFnRunner<>( - fnRunner, - ByteArrayCoder.of(), - new StreamingSideInputFetcher<>( - sideInputViews, - kwiCoder.getElementCoder(), - processFn.getInputWindowingStrategy(), - (StreamingModeExecutionContext.StreamingModeStepContext) userStepContext), - userStepContext); - } return fnRunner; } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkKitemSideInputParDoFn.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkKitemSideInputParDoFn.java new file mode 100644 index 000000000000..e789415ecbc7 --- /dev/null +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkKitemSideInputParDoFn.java @@ -0,0 +1,251 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.dataflow.worker; + +import com.google.api.client.util.Lists; +import com.google.common.collect.Iterables; +import java.io.Closeable; +import java.util.Collection; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; +import java.util.stream.StreamSupport; +import org.apache.beam.runners.core.KeyedWorkItem; +import org.apache.beam.runners.core.KeyedWorkItems; +import org.apache.beam.runners.core.SideInputReader; +import org.apache.beam.runners.core.StateNamespaces; +import org.apache.beam.runners.core.StateTag; +import org.apache.beam.runners.core.StateTags; +import org.apache.beam.runners.core.TimerInternals; +import org.apache.beam.runners.dataflow.worker.util.ValueInEmptyWindows; +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.coders.Coder; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.state.ValueState; +import org.apache.beam.sdk.transforms.DoFnSchemaInformation; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.DoFnInfo; +import org.apache.beam.sdk.values.PCollectionView; +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({ + "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) + "nullness" // TODO(https://github.com/apache/beam/issues/20497) +}) +public class StreamingKeyedWorkKitemSideInputParDoFn + implements ParDoFn { + private final StateTag> keyAddr; + private final Coder inputCoder; + private final SimpleParDoFnHelpers, OutputT, W> helpers; + protected @Nullable StreamingSideInputProcessor sideInputProcessor; + + StreamingKeyedWorkKitemSideInputParDoFn( + PipelineOptions options, + DoFnInstanceManager doFnInstanceManager, + SideInputReader sideInputReader, + TupleTag mainOutputTag, + Map, Integer> outputTupleTagsToReceiverIndices, + DataflowExecutionContext.DataflowStepContext stepContext, + DataflowOperationContext operationContext, + DoFnSchemaInformation doFnSchemaInformation, + Map> sideInputMapping, + DoFnRunnerFactory runnerFactory, + Coder keyCoder, + Coder inputCoder) { + helpers = + new SimpleParDoFnHelpers<>( + options, + doFnInstanceManager, + sideInputReader, + mainOutputTag, + outputTupleTagsToReceiverIndices, + stepContext, + operationContext, + doFnSchemaInformation, + sideInputMapping, + runnerFactory); + this.keyAddr = StateTags.makeSystemTagInternal(StateTags.value("key", keyCoder)); + this.inputCoder = inputCoder; + } + + ValueState keyValue() { + return helpers.stepContext.stateInternals().state(StateNamespaces.global(), keyAddr); + } + + @Override + public void startBundle(Receiver... receivers) throws Exception { + helpers.startBundle(receivers); + if (helpers.hasStreamingSideInput) { + // There is non-trivial setup that needs to be performed for watermark propagation + // even on empty bundles. + helpers.reallyStartBundle(); + onStartKey(); + } + } + + protected void onStartKey() { + if (helpers.hasStreamingSideInput) { + sideInputProcessor = + new StreamingSideInputProcessor<>( + new StreamingSideInputFetcher( + helpers.fnInfo.getSideInputViews(), + inputCoder, + (WindowingStrategy) helpers.fnInfo.getWindowingStrategy(), + (StreamingModeExecutionContext.StreamingModeStepContext) + helpers.userStepContext)); + } + + if (sideInputProcessor != null) { + boolean hasState = helpers.hasState(); + + // TODO(relax): We should be able to get this without writing it to state! + K key = keyValue().read(); + + Iterable> unblockedElements = + Lists.newArrayList(sideInputProcessor.tryUnblockElements()); + Iterable unblockedTimers = + Lists.newArrayList(sideInputProcessor.tryUnblockTimers()); + if (!Iterables.isEmpty(unblockedElements) || !Iterables.isEmpty(unblockedTimers)) { + helpers.fnRunner.processElement( + new ValueInEmptyWindows<>( + KeyedWorkItems.workItem(key, unblockedTimers, unblockedElements))); + } + + if (hasState) { + List windows = + (List) + StreamSupport.stream(unblockedElements.spliterator(), false) + .flatMap(wv -> wv.getWindows().stream()) + .collect(Collectors.toList()); + // These elements are now processed. Register cleanup timers for all the unblocked + // windows. + helpers.registerStateCleanup( + (WindowingStrategy) getDoFnInfo().getWindowingStrategy(), windows); + } + } + } + + @Override + @SuppressWarnings("unchecked") + public void processElement(Object untypedElem) throws Exception { + if (helpers.fnRunner == null) { + // If we need to run reallyStartBundle in here, we need to make sure to switch the state + // sampler into the start state. + try (Closeable start = helpers.operationContext.enterStart()) { + helpers.reallyStartBundle(); + onStartKey(); + } + } + helpers.outputsPerElementTracker.onProcessElement(); + + WindowedValue> elem = + (WindowedValue>) untypedElem; + onProcessWindowedValue(elem); + + helpers.outputsPerElementTracker.onProcessElementSuccess(); + } + + @Override + 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. + Coder windowCoder = + (Coder) + (helpers.fnInfo != null ? helpers.fnInfo : helpers.doFnInstanceManager.peek()) + .getWindowingStrategy() + .getWindowFn() + .windowCoder(); + // TODO: WE SHOULD FAIL HERE + helpers.processTimers( + SimpleParDoFnHelpers.TimerType.USER, + helpers.userStepContext, + windowCoder, + this::onStartKey, + sideInputProcessor); + helpers.processTimers( + SimpleParDoFnHelpers.TimerType.SYSTEM, + helpers.stepContext, + windowCoder, + this::onStartKey, + sideInputProcessor); + } + + @Override + public void finishBundle() throws Exception { + helpers.finishBundle(sideInputProcessor); + this.sideInputProcessor = null; + } + + @Override + public void abort() throws Exception { + helpers.abort(); + } + + protected void onProcessWindowedValue(WindowedValue> elem) { + // TODO: Get rid of this! + final K key = elem.getValue().key(); + keyValue().write(key); + + boolean hasState = helpers.hasState(); + Collection windowsProcessed; + if (sideInputProcessor != null) { + windowsProcessed = hasState ? Lists.newArrayList() : Collections.emptyList(); + KeyedWorkItem unblocked = sideInputProcessor.handleProcessKeyedWorkItem(elem); + if (!Iterables.isEmpty(unblocked.elementsIterable()) + || !Iterables.isEmpty(unblocked.timersIterable())) { + helpers.fnRunner.processElement(elem.withValue(unblocked)); + } + if (hasState) { + windowsProcessed = + (Collection) + StreamSupport.stream(unblocked.elementsIterable().spliterator(), false) + .flatMap(wv -> wv.getWindows().stream()) + .collect(Collectors.toList()); + } + } else { + helpers.fnRunner.processElement(elem); + windowsProcessed = (Collection) elem.getWindows(); + } + if (hasState) { + helpers.registerStateCleanup( + (WindowingStrategy) getDoFnInfo().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; + } +} diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputFetcher.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputFetcher.java index e97e16ca3133..0369b82be730 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputFetcher.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputFetcher.java @@ -332,7 +332,7 @@ private Windmill.GlobalDataRequest buildGlob .build(); } - private static class GlobalDataRequestCoder extends AtomicCoder { + static class GlobalDataRequestCoder extends AtomicCoder { private final Class protoMessageClass = Windmill.GlobalDataRequest.class; private transient Parser memoizedParser; diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputProcessor.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputProcessor.java index f51312a6e9dc..01ca0a77fd4e 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputProcessor.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputProcessor.java @@ -18,17 +18,26 @@ package org.apache.beam.runners.dataflow.worker; import java.util.Iterator; +import java.util.List; import java.util.NoSuchElementException; import java.util.Set; +import org.apache.beam.runners.core.KeyedWorkItem; +import org.apache.beam.runners.core.KeyedWorkItems; import org.apache.beam.runners.core.TimerInternals; import org.apache.beam.sdk.state.BagState; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterators; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; import org.checkerframework.checker.nullness.qual.Nullable; /** Helper class for handling elements blocked on side inputs. */ +@SuppressWarnings({ + "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) + "nullness" // TODO(https://github.com/apache/beam/issues/20497) +}) class StreamingSideInputProcessor { private final StreamingSideInputFetcher sideInputFetcher; @@ -49,62 +58,27 @@ Iterator> tryUnblockElements() { Iterable>> elementsBags = sideInputFetcher.prefetchElements(readyWindows); - // Return a lazy iterator to the released elements. This is a destructive iterator - it clears - // the bags after reading them. Bags can be paged in from the service, so we try to avoid - // materializing the whole - // bag into memory here. + // Return a lazy iterator to the released elements. Iterator> releasedElements = - new Iterator>() { - Iterator>> bagsIterator = elementsBags.iterator(); - @Nullable Iterator> currentBagElements; - @Nullable BagState> currentBag; - - @Override - public boolean hasNext() { - do { - if (currentBagElements == null || !currentBagElements.hasNext()) { - if (!advanceBag()) { - // We're done iterating - release the blocked windows. - sideInputFetcher.releaseBlockedWindows(readyWindows); - return false; - } - } - } while (!org.apache.beam.sdk.util.Preconditions.checkStateNotNull(currentBagElements) - .hasNext()); - return true; - } + new DestructivePagingIterator(elementsBags, readyWindows); - boolean advanceBag() { - // Once we finish reading a bag, clear it. - clearCurrentBag(); - if (bagsIterator.hasNext()) { - currentBag = bagsIterator.next(); - currentBagElements = currentBag.read().iterator(); - return true; - } else { - return false; - } - } + return releasedElements; + } - void clearCurrentBag() { - if (currentBag != null) { - currentBag.clear(); - currentBag = null; - currentBagElements = null; - } - } + Iterator tryUnblockTimers() { + sideInputFetcher.prefetchBlockedMap(); - @Override - public WindowedValue next() { - if (!hasNext()) { - throw new NoSuchElementException(); - } - return org.apache.beam.sdk.util.Preconditions.checkStateNotNull(currentBagElements) - .next(); - } - }; + // Find the set of ready windows. + Set readyWindows = sideInputFetcher.getReadyWindows(); - return releasedElements; + Iterable> timerBags = + sideInputFetcher.prefetchTimers(readyWindows); + + // Return a lazy iterator to the released timers. + Iterator releasedTimers = + new DestructivePagingIterator(timerBags, readyWindows); + + return releasedTimers; } void handleFinishBundle() { @@ -124,6 +98,24 @@ Iterator> handleProcessElement( (WindowedValue e) -> !sideInputFetcher.storeIfBlocked(e)); } + KeyedWorkItem handleProcessKeyedWorkItem( + WindowedValue> elem) { + List> readyInputs = + Lists.newArrayList( + Iterables.filter( + elem.getValue().elementsIterable(), + input -> !sideInputFetcher.storeIfBlocked(input))); + + List readyTimers = + Lists.newArrayList( + Iterables.filter( + elem.getValue().timersIterable(), + timer -> !sideInputFetcher.storeIfBlocked(timer))); + KeyedWorkItem keyedWorkItem = + KeyedWorkItems.workItem(elem.getValue().key(), readyTimers, readyInputs); + return keyedWorkItem; + } + void handleProcessTimer(TimerInternals.TimerData timer) { // We must call this to ensure the side-input is cached for the timer. However since a user // timer can only @@ -132,4 +124,63 @@ void handleProcessTimer(TimerInternals.TimerData timer) { // we get here. Preconditions.checkState(!sideInputFetcher.storeIfBlocked(timer)); } + + // This is a destructive iterator - it clears + // the bags after reading them. Bags can be paged in from the service, so we try to avoid + // materializing the whole + // bag into memory here. + private class DestructivePagingIterator implements Iterator { + private final Set readyWindows; + Iterator> bagsIterator; + @Nullable Iterator currentBagElements; + @Nullable BagState currentBag; + + public DestructivePagingIterator(Iterable> elementsBags, Set readyWindows) { + this.readyWindows = readyWindows; + bagsIterator = elementsBags.iterator(); + } + + @Override + public boolean hasNext() { + do { + if (currentBagElements == null || !currentBagElements.hasNext()) { + if (!advanceBag()) { + // We're done iterating - release the blocked windows. + sideInputFetcher.releaseBlockedWindows(readyWindows); + return false; + } + } + } while (!org.apache.beam.sdk.util.Preconditions.checkStateNotNull(currentBagElements) + .hasNext()); + return true; + } + + boolean advanceBag() { + // Once we finish reading a bag, clear it. + clearCurrentBag(); + if (bagsIterator.hasNext()) { + currentBag = bagsIterator.next(); + currentBagElements = currentBag.read().iterator(); + return true; + } else { + return false; + } + } + + void clearCurrentBag() { + if (currentBag != null) { + currentBag.clear(); + currentBag = null; + currentBagElements = null; + } + } + + @Override + public T next() { + if (!hasNext()) { + throw new NoSuchElementException(); + } + return org.apache.beam.sdk.util.Preconditions.checkStateNotNull(currentBagElements).next(); + } + } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UserParDoFnFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UserParDoFnFactory.java index ffeb2aa35460..6807a2ecbd5c 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UserParDoFnFactory.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UserParDoFnFactory.java @@ -24,18 +24,22 @@ import com.google.api.services.dataflow.model.SideInputInfo; import java.util.List; import java.util.Map; +import org.apache.beam.runners.core.KeyedWorkItemCoder; import org.apache.beam.runners.core.SideInputReader; import org.apache.beam.runners.dataflow.BatchStatefulParDoOverrides; import org.apache.beam.runners.dataflow.DataflowRunner; import org.apache.beam.runners.dataflow.util.CloudObject; import org.apache.beam.runners.dataflow.util.PropertyNames; import org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoFn; +import org.apache.beam.sdk.coders.ByteArrayCoder; import org.apache.beam.sdk.coders.Coder; 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.windowing.BoundedWindow; import org.apache.beam.sdk.util.DoFnInfo; import org.apache.beam.sdk.util.SerializableUtils; +import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollectionView; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.Cache; @@ -52,7 +56,7 @@ }) class UserParDoFnFactory implements ParDoFnFactory { static UserParDoFnFactory createDefault() { - return new UserParDoFnFactory(new UserDoFnExtractor(), SimpleDoFnRunnerFactory.INSTANCE_DONT_DELEGATE_SIDE_INPUTS); + return new UserParDoFnFactory(new UserDoFnExtractor(), SimpleDoFnRunnerFactory.INSTANCE, false); } interface DoFnExtractor { @@ -74,10 +78,15 @@ private static class UserDoFnExtractor implements DoFnExtractor { private final DoFnExtractor doFnExtractor; private final DoFnRunnerFactory runnerFactory; + private final boolean streamingKeyedWorkItem; - UserParDoFnFactory(DoFnExtractor doFnExtractor, DoFnRunnerFactory runnerFactory) { + UserParDoFnFactory( + DoFnExtractor doFnExtractor, + DoFnRunnerFactory runnerFactory, + boolean streamingKeyedWorkItem) { this.doFnExtractor = doFnExtractor; this.runnerFactory = runnerFactory; + this.streamingKeyedWorkItem = streamingKeyedWorkItem; } @Override @@ -144,17 +153,38 @@ public ParDoFn create( writerFn.getDataCoder(), (Coder) doFnInfo.getWindowingStrategy().getWindowFn().windowCoder()); } else { - return new SimpleParDoFn( - options, - instanceManager, - sideInputReader, - doFnInfo.getMainOutput(), - outputTupleTagsToReceiverIndices, - stepContext, - operationContext, - doFnInfo.getDoFnSchemaInformation(), - doFnInfo.getSideInputMapping(), - runnerFactory); + boolean hasStreamingSideInput = + options.as(StreamingOptions.class).isStreaming() && !sideInputReader.isEmpty(); + + if (streamingKeyedWorkItem && hasStreamingSideInput) { + KeyedWorkItemCoder> kwiCoder = + (KeyedWorkItemCoder>) doFnInfo.getInputCoder(); + return new StreamingKeyedWorkKitemSideInputParDoFn<>( + options, + instanceManager, + sideInputReader, + doFnInfo.getMainOutput(), + outputTupleTagsToReceiverIndices, + stepContext, + operationContext, + doFnInfo.getDoFnSchemaInformation(), + doFnInfo.getSideInputMapping(), + runnerFactory, + ByteArrayCoder.of(), + kwiCoder.getElementCoder()); + } else { + return new SimpleParDoFn( + options, + instanceManager, + sideInputReader, + doFnInfo.getMainOutput(), + outputTupleTagsToReceiverIndices, + stepContext, + operationContext, + doFnInfo.getDoFnSchemaInformation(), + doFnInfo.getSideInputMapping(), + runnerFactory); + } } } } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFnTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFnTest.java index 9c9f5386f443..bd742226c804 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFnTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFnTest.java @@ -92,7 +92,7 @@ public void setUp() { // TODO: Remove once Distributions has shipped. options .as(DataflowPipelineDebugOptions.class) - .setExperiments(Lists.newArrayList(SimpleParDoFn.OUTPUTS_PER_ELEMENT_EXPERIMENT)); + .setExperiments(Lists.newArrayList(SimpleParDoFnHelpers.OUTPUTS_PER_ELEMENT_EXPERIMENT)); operationContext = TestOperationContext.create(); stepContext = @@ -558,7 +558,7 @@ public void testOutputsPerElementCounter() throws Exception { public void testOutputsPerElementCounterDisabledViaExperiment() throws Exception { DataflowPipelineDebugOptions debugOptions = options.as(DataflowPipelineDebugOptions.class); List experiments = debugOptions.getExperiments(); - experiments.remove(SimpleParDoFn.OUTPUTS_PER_ELEMENT_EXPERIMENT); + experiments.remove(SimpleParDoFnHelpers.OUTPUTS_PER_ELEMENT_EXPERIMENT); debugOptions.setExperiments(experiments); List counterUpdates = executeParDoFnCounterTest(0); diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkKitemSideInputParDoFnTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkKitemSideInputParDoFnTest.java new file mode 100644 index 000000000000..8dc443f51443 --- /dev/null +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkKitemSideInputParDoFnTest.java @@ -0,0 +1,494 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.dataflow.worker; + +import static org.junit.Assert.assertEquals; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.when; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import org.apache.beam.runners.core.DoFnRunner; +import org.apache.beam.runners.core.InMemoryStateInternals; +import org.apache.beam.runners.core.KeyedWorkItem; +import org.apache.beam.runners.core.KeyedWorkItems; +import org.apache.beam.runners.core.OutputAndTimeBoundedSplittableProcessElementInvoker; +import org.apache.beam.runners.core.SideInputReader; +import org.apache.beam.runners.core.SimpleDoFnRunner; +import org.apache.beam.runners.core.SplittableParDoViaKeyedWorkItems.ProcessFn; +import org.apache.beam.runners.core.StateInternals; +import org.apache.beam.runners.core.StateNamespaces; +import org.apache.beam.runners.core.TimerInternals; +import org.apache.beam.runners.core.TimerInternals.TimerData; +import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputState; +import org.apache.beam.runners.dataflow.worker.util.ValueInEmptyWindows; +import org.apache.beam.runners.dataflow.worker.util.common.worker.Receiver; +import org.apache.beam.runners.dataflow.worker.windmill.Windmill; +import org.apache.beam.runners.dataflow.worker.windmill.Windmill.Timer; +import org.apache.beam.sdk.coders.BigEndianIntegerCoder; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.options.StreamingOptions; +import org.apache.beam.sdk.state.TimeDomain; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.DoFnSchemaInformation; +import org.apache.beam.sdk.transforms.View; +import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker; +import org.apache.beam.sdk.transforms.splittabledofn.SplitResult; +import org.apache.beam.sdk.transforms.windowing.FixedWindows; +import org.apache.beam.sdk.transforms.windowing.IntervalWindow; +import org.apache.beam.sdk.transforms.windowing.PaneInfo; +import org.apache.beam.sdk.transforms.windowing.Window; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.util.DoFnInfo; +import org.apache.beam.sdk.util.WindowedValueMultiReceiver; +import org.apache.beam.sdk.values.CausedByDrain; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PCollectionView; +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.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.joda.time.Duration; +import org.joda.time.Instant; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; + +/** Unit tests for {@link StreamingKeyedWorkKitemSideInputParDoFn}. */ +@RunWith(JUnit4.class) +public class StreamingKeyedWorkKitemSideInputParDoFnTest { + private static final FixedWindows WINDOW_FN = FixedWindows.of(Duration.millis(10)); + private static TupleTag> mainOutputTag = new TupleTag<>(); + + private final InMemoryStateInternals state = InMemoryStateInternals.forKey("a"); + + @Mock private StreamingModeExecutionContext.StepContext stepContext; + @Mock private TimerInternals mockTimerInternals; + @Mock private SideInputReader mockSideInputReader; + + @SuppressWarnings({"unchecked", "rawtypes"}) + @Before + public void setUp() { + MockitoAnnotations.initMocks(this); + when(stepContext.stateInternals()).thenReturn((StateInternals) state); + when(stepContext.timerInternals()).thenReturn(mockTimerInternals); + when(stepContext.namespacedToUser()).thenReturn(stepContext); + when(mockSideInputReader.isEmpty()).thenReturn(false); + } + + @Test + public void testInvokeProcessElement() throws Exception { + PCollectionView view = createView(); + + when(stepContext.issueSideInputFetch( + eq(view), + any(org.apache.beam.sdk.transforms.windowing.BoundedWindow.class), + eq(SideInputState.UNKNOWN))) + .thenReturn(false); + when(stepContext.issueSideInputFetch( + eq(view), + any(org.apache.beam.sdk.transforms.windowing.BoundedWindow.class), + eq(SideInputState.KNOWN_READY))) + .thenReturn(true); + + when(mockTimerInternals.currentInputWatermarkTime()).thenReturn(new Instant(15L)); + StreamingKeyedWorkKitemSideInputParDoFn, IntervalWindow> + runner = createRunner(view); + + TestReceiver receiver = new TestReceiver(); + runner.startBundle(receiver); + + KeyedWorkItem elemsWorkItem = + KeyedWorkItems.elementsWorkItem( + "a", + ImmutableList.of( + createDatum(13, 13L), + createDatum(16, 16L), // side inputs non-ready element + createDatum(18, 18L))); + + runner.processElement(new ValueInEmptyWindows<>(elemsWorkItem)); + + // Initially blocked! No output. + assertEquals(0, receiver.outputs.size()); + + when(mockTimerInternals.currentInputWatermarkTime()).thenReturn(new Instant(20)); + runner.processElement( + new ValueInEmptyWindows<>( + KeyedWorkItems.timersWorkItem( + "a", + ImmutableList.of( + timerData(window(10, 20), new Instant(19), Timer.Type.WATERMARK))))); + + // Timer is blocked too! + assertEquals(0, receiver.outputs.size()); + + // Now make it ready! + IntervalWindow readyWindow = window(10, 20); + Windmill.GlobalDataId id = + Windmill.GlobalDataId.newBuilder() + .setTag(view.getTagInternal().getId()) + .setVersion( + ByteString.copyFrom( + CoderUtils.encodeToByteArray(IntervalWindow.getCoder(), readyWindow))) + .build(); + + when(stepContext.getSideInputNotifications()) + .thenReturn(Arrays.asList(id)); + + runner.finishBundle(); + + runner.startBundle(receiver); + + // We don't check for output here because we just wanted to see if the runner works + // without exceptions. The issue was lifecycle of the runner bundle (finishBundle, startBundle). + } + + static class TestSplittableDoFn extends DoFn> { + @ProcessElement + public void processElement(ProcessContext c, RestrictionTracker tracker) { + c.output(KV.of(tracker.currentRestriction(), c.element())); + } + + @GetInitialRestriction + public String getInitialRestriction(@Element Integer element) { + return "restriction"; + } + + @NewTracker + public RestrictionTracker newTracker(@Restriction String restriction) { + return new RestrictionTracker() { + @Override + public boolean tryClaim(Object position) { + return true; + } + + @Override + public String currentRestriction() { + return restriction; + } + + @Override + public SplitResult trySplit(double fractionOfRemainder) { + return null; + } + + @Override + public void checkDone() {} + + @Override + public IsBounded isBounded() { + return IsBounded.BOUNDED; + } + }; + } + } + + @Test + public void testSplittableProcessElement() throws Exception { + PCollectionView view = createView(); + + when(stepContext.issueSideInputFetch( + eq(view), + any(org.apache.beam.sdk.transforms.windowing.BoundedWindow.class), + eq(SideInputState.UNKNOWN))) + .thenReturn(false); + when(stepContext.issueSideInputFetch( + eq(view), + any(org.apache.beam.sdk.transforms.windowing.BoundedWindow.class), + eq(SideInputState.KNOWN_READY))) + .thenReturn(true); + + when(mockTimerInternals.currentInputWatermarkTime()).thenReturn(new Instant(15L)); + StreamingKeyedWorkKitemSideInputParDoFn< + byte[], KV, KV, IntervalWindow> + runner = createSplittableRunner(view); + + TestReceiver receiver = new TestReceiver(); + runner.startBundle(receiver); + + KeyedWorkItem> elemsWorkItem = + KeyedWorkItems.elementsWorkItem( + new byte[] {1}, ImmutableList.of(createDatum(KV.of(13, "restriction"), 13L))); + + runner.processElement(new ValueInEmptyWindows<>(elemsWorkItem)); + + // Initially blocked! No output. + assertEquals(0, receiver.outputs.size()); + runner.finishBundle(); + + // Now make it ready! + IntervalWindow readyWindow = window(10, 20); + Windmill.GlobalDataId id = + Windmill.GlobalDataId.newBuilder() + .setTag(view.getTagInternal().getId()) + .setVersion( + ByteString.copyFrom( + CoderUtils.encodeToByteArray(IntervalWindow.getCoder(), readyWindow))) + .build(); + + when(stepContext.getSideInputNotifications()) + .thenReturn(Arrays.asList(id)); + + runner.startBundle(receiver); + + // Note: unblocking logic would run here if the environment is fully mocked to push + // blocked items back into processing. For the purpose of testing SplittableDoFn initialization, + // this suffices. + } + + private WindowedValue createDatum(T element, long timestampMillis) { + Instant timestamp = new Instant(timestampMillis); + return WindowedValues.of( + element, timestamp, Arrays.asList(WINDOW_FN.assignWindow(timestamp)), PaneInfo.NO_FIRING); + } + + private TimerData timerData(IntervalWindow window, Instant timestamp, Timer.Type type) { + return TimerData.of( + StateNamespaces.window(IntervalWindow.getCoder(), window), + timestamp, + timestamp, + type == Windmill.Timer.Type.WATERMARK ? TimeDomain.EVENT_TIME : TimeDomain.PROCESSING_TIME, + CausedByDrain.NORMAL); + } + + private IntervalWindow window(long start, long end) { + return new IntervalWindow(new Instant(start), new Instant(end)); + } + + private PCollectionView createView() { + return TestPipeline.create() + .apply(Create.empty(StringUtf8Coder.of())) + .apply(Window.into(WINDOW_FN)) + .apply(View.asSingleton()); + } + + static class TestReceiver implements Receiver { + List>> outputs = new ArrayList<>(); + + @SuppressWarnings("unchecked") + @Override + public void process(Object outputElem) { + outputs.add((WindowedValue>) outputElem); + } + } + + @SuppressWarnings("unchecked") + private StreamingKeyedWorkKitemSideInputParDoFn< + String, Integer, KV, IntervalWindow> + createRunner(PCollectionView view) throws Exception { + Coder keyCoder = StringUtf8Coder.of(); + Coder inputCoder = BigEndianIntegerCoder.of(); + + WindowingStrategy windowingStrategy = WindowingStrategy.of(WINDOW_FN); + + DoFn, KV> theDoFn = + new DoFn, KV>() { + @ProcessElement + public void processElement(ProcessContext c) { + KeyedWorkItem kwi = c.element(); + for (WindowedValue wv : kwi.elementsIterable()) { + c.output(KV.of(kwi.key(), wv.getValue())); + } + } + }; + + DoFnInfo, KV> fnInfo = + DoFnInfo.forFn( + theDoFn, + windowingStrategy, + ImmutableList.of(view), + (Coder) null, + mainOutputTag, + DoFnSchemaInformation.create(), + Collections.emptyMap()); + + DoFnRunnerFactory, KV> runnerFactory = + new DoFnRunnerFactory, KV>() { + @Override + public DoFnRunner, KV> createRunner( + DoFn, KV> fn, + PipelineOptions options, + TupleTag> mainOutputTag, + List> sideOutputTags, + Iterable> sideInputViews, + SideInputReader sideInputReader, + Coder> inputCoder, + Map, Coder> outputCoders, + WindowingStrategy windowingStrategy, + DataflowExecutionContext.DataflowStepContext stepContext, + DataflowExecutionContext.DataflowStepContext userStepContext, + WindowedValueMultiReceiver outputManager2, + DoFnSchemaInformation doFnSchemaInformation, + Map> sideInputMapping) { + return new SimpleDoFnRunner<>( + options, + fn, + sideInputReader, + outputManager2, + mainOutputTag, + sideOutputTags, + stepContext, + inputCoder, + outputCoders, + windowingStrategy, + doFnSchemaInformation, + sideInputMapping); + } + }; + + PipelineOptions options = PipelineOptionsFactory.create(); + options.as(StreamingOptions.class).setStreaming(true); + + return new StreamingKeyedWorkKitemSideInputParDoFn<>( + options, + DoFnInstanceManagers.singleInstance(fnInfo), + mockSideInputReader, + mainOutputTag, + ImmutableMap.of(mainOutputTag, 0), + stepContext, + TestOperationContext.create(), + DoFnSchemaInformation.create(), + Collections.emptyMap(), + runnerFactory, + keyCoder, + inputCoder); + } + + @SuppressWarnings("unchecked") + private StreamingKeyedWorkKitemSideInputParDoFn< + byte[], KV, KV, IntervalWindow> + createSplittableRunner(PCollectionView view) throws Exception { + org.apache.beam.sdk.coders.ByteArrayCoder keyCoder = + org.apache.beam.sdk.coders.ByteArrayCoder.of(); + Coder> inputCoder = + org.apache.beam.sdk.coders.KvCoder.of(BigEndianIntegerCoder.of(), StringUtf8Coder.of()); + + WindowingStrategy windowingStrategy = + (WindowingStrategy) WindowingStrategy.of(WINDOW_FN); + + TestSplittableDoFn theDoFn = new TestSplittableDoFn(); + + ProcessFn, String, Object, Object> processFn = + new ProcessFn, String, Object, Object>( + theDoFn, + BigEndianIntegerCoder.of(), + StringUtf8Coder.of(), + (Coder) StringUtf8Coder.of(), // watermarkEstimatorStateCoder + windowingStrategy, + Collections.emptyMap()); + processFn.setup(org.apache.beam.sdk.options.PipelineOptionsFactory.create()); + + DoFnInfo>, KV> fnInfo = + DoFnInfo.forFn( + processFn, + windowingStrategy, + ImmutableList.of(view), + (Coder) null, + mainOutputTag, + DoFnSchemaInformation.create(), + Collections.emptyMap()); + + DoFnRunnerFactory>, KV> + runnerFactory = + new DoFnRunnerFactory< + KeyedWorkItem>, KV>() { + @Override + public DoFnRunner>, KV> + createRunner( + DoFn>, KV> fn, + PipelineOptions options, + TupleTag> mainOutputTag, + List> sideOutputTags, + Iterable> sideInputViews, + SideInputReader sideInputReader, + Coder>> inputCoder, + Map, Coder> outputCoders, + WindowingStrategy windowingStrategy, + DataflowExecutionContext.DataflowStepContext stepContext, + DataflowExecutionContext.DataflowStepContext userStepContext, + WindowedValueMultiReceiver outputManager2, + DoFnSchemaInformation doFnSchemaInformation, + Map> sideInputMapping) { + + ProcessFn, String, Object, Object> fn2 = + (ProcessFn, String, Object, Object>) fn; + fn2.setStateInternalsFactory(key -> (StateInternals) stepContext.stateInternals()); + fn2.setTimerInternalsFactory(key -> stepContext.timerInternals()); + fn2.setSideInputReader(sideInputReader); + fn2.setProcessElementInvoker( + new OutputAndTimeBoundedSplittableProcessElementInvoker< + Integer, KV, String, Object, Object>( + fn2.getFn(), + options, + outputManager2, + mainOutputTag, + sideInputReader, + java.util.concurrent.Executors.newSingleThreadScheduledExecutor(), + 10000, + Duration.standardSeconds(10), + () -> null)); + + return new SimpleDoFnRunner<>( + options, + fn, + sideInputReader, + outputManager2, + mainOutputTag, + sideOutputTags, + stepContext, + inputCoder, + outputCoders, + windowingStrategy, + doFnSchemaInformation, + sideInputMapping); + } + }; + + PipelineOptions options = PipelineOptionsFactory.create(); + options.as(StreamingOptions.class).setStreaming(true); + + return new StreamingKeyedWorkKitemSideInputParDoFn<>( + options, + DoFnInstanceManagers.singleInstance(fnInfo), + mockSideInputReader, + mainOutputTag, + ImmutableMap.of(mainOutputTag, 0), + stepContext, + TestOperationContext.create(), + DoFnSchemaInformation.create(), + Collections.emptyMap(), + runnerFactory, + keyCoder, + inputCoder); + } +} 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 9d3fa9b211b1..6bb0203231a0 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 @@ -383,7 +383,7 @@ public void testCleanupRegistered() throws Exception { verify(stepContext) .setStateCleanupTimer( - SimpleParDoFn.CLEANUP_TIMER_ID, + SimpleParDoFnHelpers.CLEANUP_TIMER_ID, firstWindow, IntervalWindow.getCoder(), firstWindow.maxTimestamp().plus(Duration.millis(1L)), @@ -443,7 +443,7 @@ public void testCleanupTimerForGlobalWindowWithAllowedLateness() throws Exceptio greaterThan(BoundedWindow.TIMESTAMP_MAX_VALUE)); verify(stepContext) .setStateCleanupTimer( - SimpleParDoFn.CLEANUP_TIMER_ID, + SimpleParDoFnHelpers.CLEANUP_TIMER_ID, globalWindow, GlobalWindow.Coder.INSTANCE, BoundedWindow.TIMESTAMP_MAX_VALUE, @@ -459,7 +459,7 @@ public void testCleanupTimerForGlobalWindowWithAllowedLateness() throws Exceptio when(stepContext.getNextFiredTimer((Coder) GlobalWindow.Coder.INSTANCE)) .thenReturn( TimerData.of( - SimpleParDoFn.CLEANUP_TIMER_ID, + SimpleParDoFnHelpers.CLEANUP_TIMER_ID, globalWindowNamespace, BoundedWindow.TIMESTAMP_MAX_VALUE, BoundedWindow.TIMESTAMP_MAX_VALUE.minus(Duration.millis(1)), @@ -535,7 +535,7 @@ public void testCleanupWorks() throws Exception { when(stepContext.getNextFiredTimer(windowCoder)) .thenReturn( TimerData.of( - SimpleParDoFn.CLEANUP_TIMER_ID, + SimpleParDoFnHelpers.CLEANUP_TIMER_ID, firstWindowNamespace, firstWindow.maxTimestamp().plus(Duration.millis(1L)), firstWindow.maxTimestamp().plus(Duration.millis(1L)), @@ -552,7 +552,7 @@ public void testCleanupWorks() throws Exception { when(stepContext.getNextFiredTimer((Coder) windowCoder)) .thenReturn( TimerData.of( - SimpleParDoFn.CLEANUP_TIMER_ID, + SimpleParDoFnHelpers.CLEANUP_TIMER_ID, secondWindowNamespace, secondWindow.maxTimestamp().plus(Duration.millis(1L)), secondWindow.maxTimestamp().plus(Duration.millis(1L)), From 5e402b82304c682272710f42f35155a0a1c78b58 Mon Sep 17 00:00:00 2001 From: Reuven Lax Date: Wed, 3 Jun 2026 20:30:24 -0700 Subject: [PATCH 06/11] fix splittable dofn --- .../dataflow/worker/SimpleParDoFn.java | 25 ++-- .../dataflow/worker/SimpleParDoFnHelpers.java | 2 +- ...reamingKeyedWorkKitemSideInputParDoFn.java | 42 +++--- .../worker/StreamingSideInputDoFnRunner.java | 11 +- .../worker/StreamingSideInputProcessor.java | 97 +++---------- .../worker/SimpleParDoFnHelpersTest.java | 135 ++++++++++++++++++ ...gKeyedWorkItemSideInputDoFnRunnerTest.java | 3 +- .../StreamingSideInputProcessorTest.java | 12 +- .../sdk/transforms/SplittableDoFnTest.java | 2 +- 9 files changed, 203 insertions(+), 126 deletions(-) create mode 100644 runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFnHelpersTest.java 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 56f47dbcf067..90492d4ee8d3 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 @@ -104,18 +104,19 @@ protected void onStartKey() { helpers.userStepContext)); boolean hasState = helpers.hasState(); - Iterator> unblockedElements = sideInputProcessor.tryUnblockElements(); - for (Iterator> it = unblockedElements; it.hasNext(); ) { - WindowedValue unblockedElement = it.next(); - helpers.fnRunner.processElement(unblockedElement); - if (hasState) { - // These elements are now processed. Register cleanup timers for all the unblocked - // windows. - helpers.registerStateCleanup( - (WindowingStrategy) getDoFnInfo().getWindowingStrategy(), - (Collection) unblockedElement.getWindows()); - } - } + sideInputProcessor.tryUnblockElements( + unblockedElements -> { + for (WindowedValue unblockedElement : unblockedElements) { + helpers.fnRunner.processElement(unblockedElement); + if (hasState) { + // These elements are now processed. Register cleanup timers for all the unblocked + // windows. + helpers.registerStateCleanup( + (WindowingStrategy) getDoFnInfo().getWindowingStrategy(), + (Collection) unblockedElement.getWindows()); + } + } + }); } } 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 6a969c97b06a..b765228fb1fd 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 @@ -67,7 +67,7 @@ "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) "nullness" // TODO(https://github.com/apache/beam/issues/20497) }) -public class SimpleParDoFnHelpers { +class SimpleParDoFnHelpers { private static final Logger LOG = LoggerFactory.getLogger(SimpleParDoFnHelpers.class); // TODO: Remove once Distributions has shipped. diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkKitemSideInputParDoFn.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkKitemSideInputParDoFn.java index e789415ecbc7..28e27d6453e9 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkKitemSideInputParDoFn.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkKitemSideInputParDoFn.java @@ -32,7 +32,6 @@ import org.apache.beam.runners.core.StateNamespaces; import org.apache.beam.runners.core.StateTag; import org.apache.beam.runners.core.StateTags; -import org.apache.beam.runners.core.TimerInternals; import org.apache.beam.runners.dataflow.worker.util.ValueInEmptyWindows; import org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoFn; import org.apache.beam.runners.dataflow.worker.util.common.worker.Receiver; @@ -53,6 +52,7 @@ "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) "nullness" // TODO(https://github.com/apache/beam/issues/20497) }) +/** Similar to {@link SimpleParDoFn} but for splittable ProcessFns. */ public class StreamingKeyedWorkKitemSideInputParDoFn implements ParDoFn { private final StateTag> keyAddr; @@ -122,27 +122,25 @@ protected void onStartKey() { // TODO(relax): We should be able to get this without writing it to state! K key = keyValue().read(); - Iterable> unblockedElements = - Lists.newArrayList(sideInputProcessor.tryUnblockElements()); - Iterable unblockedTimers = - Lists.newArrayList(sideInputProcessor.tryUnblockTimers()); - if (!Iterables.isEmpty(unblockedElements) || !Iterables.isEmpty(unblockedTimers)) { - helpers.fnRunner.processElement( - new ValueInEmptyWindows<>( - KeyedWorkItems.workItem(key, unblockedTimers, unblockedElements))); - } - - if (hasState) { - List windows = - (List) - StreamSupport.stream(unblockedElements.spliterator(), false) - .flatMap(wv -> wv.getWindows().stream()) - .collect(Collectors.toList()); - // These elements are now processed. Register cleanup timers for all the unblocked - // windows. - helpers.registerStateCleanup( - (WindowingStrategy) getDoFnInfo().getWindowingStrategy(), windows); - } + sideInputProcessor.tryUnblockElementsAndTimers( + (unblockedElements, unblockedTimers) -> { + if (!Iterables.isEmpty(unblockedElements) || !Iterables.isEmpty(unblockedTimers)) { + helpers.fnRunner.processElement( + new ValueInEmptyWindows<>( + KeyedWorkItems.workItem(key, unblockedTimers, unblockedElements))); + } + if (hasState) { + List windows = + (List) + StreamSupport.stream(unblockedElements.spliterator(), false) + .flatMap(wv -> wv.getWindows().stream()) + .collect(Collectors.toList()); + // These elements are now processed. Register cleanup timers for all the unblocked + // windows. + helpers.registerStateCleanup( + (WindowingStrategy) getDoFnInfo().getWindowingStrategy(), windows); + } + }); } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputDoFnRunner.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputDoFnRunner.java index b41b0c5049a7..a64d1a970d34 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputDoFnRunner.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputDoFnRunner.java @@ -48,11 +48,12 @@ public StreamingSideInputDoFnRunner( @Override public void startBundle() { simpleDoFnRunner.startBundle(); - Iterator> unblocked = sideInputProcessor.tryUnblockElements(); - for (Iterator> it = unblocked; it.hasNext(); ) { - WindowedValue elem = it.next(); - simpleDoFnRunner.processElement(elem); - } + sideInputProcessor.tryUnblockElements( + unblocked -> { + for (WindowedValue elem : unblocked) { + simpleDoFnRunner.processElement(elem); + } + }); } @Override diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputProcessor.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputProcessor.java index 01ca0a77fd4e..477512f157f8 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputProcessor.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputProcessor.java @@ -19,8 +19,9 @@ import java.util.Iterator; import java.util.List; -import java.util.NoSuchElementException; import java.util.Set; +import java.util.function.BiConsumer; +import java.util.function.Consumer; import org.apache.beam.runners.core.KeyedWorkItem; import org.apache.beam.runners.core.KeyedWorkItems; import org.apache.beam.runners.core.TimerInternals; @@ -31,7 +32,6 @@ 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.collect.Iterators; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; -import org.checkerframework.checker.nullness.qual.Nullable; /** Helper class for handling elements blocked on side inputs. */ @SuppressWarnings({ @@ -49,23 +49,23 @@ public StreamingSideInputProcessor(StreamingSideInputFetcher sideInpu * Handle's startBundle. If there are unblocked elements, process them and then return the set of * windows that were unblocked. */ - Iterator> tryUnblockElements() { + void tryUnblockElements(Consumer>> consumer) { sideInputFetcher.prefetchBlockedMap(); // Find the set of ready windows. Set readyWindows = sideInputFetcher.getReadyWindows(); - Iterable>> elementsBags = + Iterable>> elementBags = sideInputFetcher.prefetchElements(readyWindows); - - // Return a lazy iterator to the released elements. - Iterator> releasedElements = - new DestructivePagingIterator(elementsBags, readyWindows); - - return releasedElements; + Iterable> releasedElements = + Iterables.concat(Iterables.transform(elementBags, BagState::read)); + consumer.accept(releasedElements); + elementBags.forEach(BagState::clear); + sideInputFetcher.releaseBlockedWindows(readyWindows); } - Iterator tryUnblockTimers() { + void tryUnblockElementsAndTimers( + BiConsumer>, Iterable> consumer) { sideInputFetcher.prefetchBlockedMap(); // Find the set of ready windows. @@ -73,12 +73,18 @@ Iterator tryUnblockTimers() { Iterable> timerBags = sideInputFetcher.prefetchTimers(readyWindows); + Iterable releasedTimers = + Iterables.concat( + Iterables.transform(sideInputFetcher.prefetchTimers(readyWindows), BagState::read)); + Iterable>> elementBags = + sideInputFetcher.prefetchElements(readyWindows); + Iterable> releasedElements = + Iterables.concat(Iterables.transform(elementBags, BagState::read)); - // Return a lazy iterator to the released timers. - Iterator releasedTimers = - new DestructivePagingIterator(timerBags, readyWindows); - - return releasedTimers; + consumer.accept(releasedElements, releasedTimers); + timerBags.forEach(BagState::clear); + elementBags.forEach(BagState::clear); + sideInputFetcher.releaseBlockedWindows(readyWindows); } void handleFinishBundle() { @@ -124,63 +130,4 @@ void handleProcessTimer(TimerInternals.TimerData timer) { // we get here. Preconditions.checkState(!sideInputFetcher.storeIfBlocked(timer)); } - - // This is a destructive iterator - it clears - // the bags after reading them. Bags can be paged in from the service, so we try to avoid - // materializing the whole - // bag into memory here. - private class DestructivePagingIterator implements Iterator { - private final Set readyWindows; - Iterator> bagsIterator; - @Nullable Iterator currentBagElements; - @Nullable BagState currentBag; - - public DestructivePagingIterator(Iterable> elementsBags, Set readyWindows) { - this.readyWindows = readyWindows; - bagsIterator = elementsBags.iterator(); - } - - @Override - public boolean hasNext() { - do { - if (currentBagElements == null || !currentBagElements.hasNext()) { - if (!advanceBag()) { - // We're done iterating - release the blocked windows. - sideInputFetcher.releaseBlockedWindows(readyWindows); - return false; - } - } - } while (!org.apache.beam.sdk.util.Preconditions.checkStateNotNull(currentBagElements) - .hasNext()); - return true; - } - - boolean advanceBag() { - // Once we finish reading a bag, clear it. - clearCurrentBag(); - if (bagsIterator.hasNext()) { - currentBag = bagsIterator.next(); - currentBagElements = currentBag.read().iterator(); - return true; - } else { - return false; - } - } - - void clearCurrentBag() { - if (currentBag != null) { - currentBag.clear(); - currentBag = null; - currentBagElements = null; - } - } - - @Override - public T next() { - if (!hasNext()) { - throw new NoSuchElementException(); - } - return org.apache.beam.sdk.util.Preconditions.checkStateNotNull(currentBagElements).next(); - } - } } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFnHelpersTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFnHelpersTest.java new file mode 100644 index 000000000000..06599870fe42 --- /dev/null +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFnHelpersTest.java @@ -0,0 +1,135 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.dataflow.worker; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import org.apache.beam.runners.core.DoFnRunner; +import org.apache.beam.runners.core.SideInputReader; +import org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.transforms.DoFnSchemaInformation; +import org.apache.beam.sdk.values.TupleTag; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; + +@RunWith(JUnit4.class) +public class SimpleParDoFnHelpersTest { + private PipelineOptions options; + @Mock DoFnInstanceManager doFnInstanceManager; + @Mock SideInputReader sideInputReader; + @Mock DataflowStepContext stepContext; + @Mock DataflowStepContext userStepContext; + @Mock DataflowOperationContext operationContext; + @Mock DoFnRunnerFactory runnerFactory; + @Mock DoFnRunner mockRunner; + + @Mock + StreamingSideInputProcessor + sideInputProcessor; + + @Mock org.apache.beam.sdk.util.DoFnInfo doFnInfo; + @Mock org.apache.beam.runners.dataflow.worker.counters.CounterFactory counterFactory; + + private static class TestDoFn extends org.apache.beam.sdk.transforms.DoFn { + @ProcessElement + public void processElement() {} + } + + private TestDoFn doFn = new TestDoFn(); + + private SimpleParDoFnHelpers< + String, String, org.apache.beam.sdk.transforms.windowing.GlobalWindow> + helpers; + + @Before + @SuppressWarnings("unchecked") + public void setUp() throws Exception { + MockitoAnnotations.initMocks(this); + options = PipelineOptionsFactory.create(); + when(stepContext.namespacedToUser()).thenReturn(userStepContext); + when(operationContext.counterFactory()).thenReturn(counterFactory); + + when(doFnInstanceManager.get()).thenReturn((org.apache.beam.sdk.util.DoFnInfo) doFnInfo); + when(doFnInfo.getDoFn()).thenReturn(doFn); + + when(runnerFactory.createRunner( + any(), any(), any(), any(), any(), any(), any(), any(), any(), any(), any(), any(), + any(), any())) + .thenReturn(mockRunner); + + helpers = + new SimpleParDoFnHelpers<>( + options, + doFnInstanceManager, + sideInputReader, + new TupleTag<>("main"), + ImmutableMap.of(new TupleTag<>("main"), 0), + stepContext, + operationContext, + DoFnSchemaInformation.create(), + ImmutableMap.of(), + runnerFactory); + } + + @Test + public void testReallyStartBundle() throws Exception { + helpers.startBundle( + mock(org.apache.beam.runners.dataflow.worker.util.common.worker.Receiver.class)); + helpers.reallyStartBundle(); + + verify(runnerFactory) + .createRunner( + any(), any(), any(), any(), any(), any(), any(), any(), any(), any(), any(), any(), + any(), any()); + verify(mockRunner).startBundle(); + } + + @Test + public void testFinishBundle() throws Exception { + helpers.startBundle( + mock(org.apache.beam.runners.dataflow.worker.util.common.worker.Receiver.class)); + helpers.reallyStartBundle(); + + helpers.finishBundle(sideInputProcessor); + + verify(mockRunner).finishBundle(); + verify(sideInputProcessor).handleFinishBundle(); + verify(doFnInstanceManager).complete(any()); + } + + @Test + public void testAbort() throws Exception { + helpers.startBundle( + mock(org.apache.beam.runners.dataflow.worker.util.common.worker.Receiver.class)); + helpers.reallyStartBundle(); + + helpers.abort(); + + verify(doFnInstanceManager).abort(any()); + } +} diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputDoFnRunnerTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputDoFnRunnerTest.java index e12ddd95f913..654707aae912 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputDoFnRunnerTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputDoFnRunnerTest.java @@ -202,8 +202,7 @@ private IntervalWindow window(long start, long end) { (WindowedValue> windowedValue) -> outputManager.output(mainOutputTag, windowedValue), stepContext); - return new StreamingKeyedWorkItemSideInputDoFnRunner< - String, Integer, KV, IntervalWindow>( + return new StreamingKeyedWorkItemSideInputDoFnRunner<>( simpleDoFnRunner, keyCoder, sideInputFetcher, stepContext); } } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputProcessorTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputProcessorTest.java index 804c22ea8567..e6e553e238ce 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputProcessorTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputProcessorTest.java @@ -19,13 +19,12 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.containsInAnyOrder; +import static org.hamcrest.Matchers.emptyIterable; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertThrows; import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anySet; import static org.mockito.Mockito.doNothing; import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -70,10 +69,9 @@ public void testTryUnblockElementsNoReadyWindows() { when(mockFetcher.getReadyWindows()).thenReturn(Collections.emptySet()); // When - Iterator> unblocked = processor.tryUnblockElements(); + processor.tryUnblockElements(unblocked -> assertThat(unblocked, emptyIterable())); // Then - assertFalse(unblocked.hasNext()); verify(mockFetcher).prefetchBlockedMap(); verify(mockFetcher).getReadyWindows(); } @@ -104,12 +102,10 @@ public void testTryUnblockElementsWithReadyWindows() { doNothing().when(mockFetcher).releaseBlockedWindows(readyWindows); // When - Iterable> unblocked = () -> processor.tryUnblockElements(); + processor.tryUnblockElements( + unblocked -> assertThat(unblocked, containsInAnyOrder(element1, element2))); // Then - verify(mockBag1, never()).clear(); - verify(mockFetcher, never()).releaseBlockedWindows(anySet()); - assertThat(unblocked, containsInAnyOrder(element1, element2)); verify(mockFetcher).prefetchBlockedMap(); verify(mockFetcher).getReadyWindows(); verify(mockFetcher).prefetchElements(readyWindows); diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/SplittableDoFnTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/SplittableDoFnTest.java index 80d8728aa01b..117ceac32d71 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/SplittableDoFnTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/SplittableDoFnTest.java @@ -426,7 +426,7 @@ private static SDFWithSideInputBase sdfWithSideInput( @Test @Category({ValidatesRunner.class, UsesBoundedSplittableParDo.class, UsesSideInputs.class}) - public void testSideInputBounded() { + public void tBounded() { testSideInput(IsBounded.BOUNDED); } From e366457229b0e83a1ddd06d6224fa57116e59206 Mon Sep 17 00:00:00 2001 From: Reuven Lax Date: Wed, 3 Jun 2026 20:47:34 -0700 Subject: [PATCH 07/11] fix --- .../runners/dataflow/worker/SimpleParDoFnHelpers.java | 11 +++++++++++ .../StreamingKeyedWorkKitemSideInputParDoFn.java | 3 +-- 2 files changed, 12 insertions(+), 2 deletions(-) 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 b765228fb1fd..5751f9619ac5 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 @@ -264,6 +264,17 @@ public void processTimer( doFn.processUserTimer(timer, sideInputProcessor); } }, + FAIL_USER { + @Override + public void processTimer( + SimpleParDoFnHelpers doFn, + TimerInternals.TimerData timer, + StreamingSideInputProcessor sideInputProcessor) + throws Exception { + throw new UnsupportedOperationException( + "Attempt to deliver a timer to a DoFn, but timers are not supported here."); + } + }, SYSTEM { @Override public void processTimer( diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkKitemSideInputParDoFn.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkKitemSideInputParDoFn.java index 28e27d6453e9..8dc6115f74f1 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkKitemSideInputParDoFn.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkKitemSideInputParDoFn.java @@ -178,9 +178,8 @@ public void processTimers() throws Exception { .getWindowingStrategy() .getWindowFn() .windowCoder(); - // TODO: WE SHOULD FAIL HERE helpers.processTimers( - SimpleParDoFnHelpers.TimerType.USER, + SimpleParDoFnHelpers.TimerType.FAIL_USER, helpers.userStepContext, windowCoder, this::onStartKey, From e71e5641638fc1ac66d6605dd1ede76ac3af938d Mon Sep 17 00:00:00 2001 From: Reuven Lax Date: Thu, 4 Jun 2026 08:50:10 -0700 Subject: [PATCH 08/11] foo --- .../dataflow/worker/SimpleParDoFnHelpers.java | 6 +- ...reamingKeyedWorkKitemSideInputParDoFn.java | 2 +- .../worker/StreamingSideInputProcessor.java | 6 +- .../worker/SimpleParDoFnHelpersTest.java | 30 +++++----- .../dataflow/worker/SimpleParDoFnTest.java | 2 +- ...ingKeyedWorkKitemSideInputParDoFnTest.java | 56 +++++++++---------- .../StreamingSideInputProcessorTest.java | 56 ++++++++++++------- .../worker/UserParDoFnFactoryTest.java | 17 +++--- 8 files changed, 89 insertions(+), 86 deletions(-) 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 5751f9619ac5..df0af65a4b9e 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 @@ -42,6 +42,7 @@ import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.StreamingOptions; +import org.apache.beam.sdk.state.State; import org.apache.beam.sdk.state.StateSpec; import org.apache.beam.sdk.state.TimeDomain; import org.apache.beam.sdk.transforms.DoFnSchemaInformation; @@ -318,8 +319,7 @@ void processTimers( } protected void processUserTimer( - TimerInternals.TimerData timer, StreamingSideInputProcessor sideInputProcessor) - throws Exception { + TimerInternals.TimerData timer, StreamingSideInputProcessor sideInputProcessor) { if (fnSignature.timerDeclarations().containsKey(timer.getTimerId()) || fnSignature.timerFamilyDeclarations().containsKey(timer.getTimerFamilyId())) { BoundedWindow window = ((StateNamespaces.WindowNamespace) timer.getNamespace()).getWindow(); @@ -393,7 +393,7 @@ private void processSystemTimer( } StateInternals stateInternals = userStepContext.stateInternals(); - org.apache.beam.sdk.state.State state = stateInternals.state(timer.getNamespace(), tag); + State state = stateInternals.state(timer.getNamespace(), tag); state.clear(); } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkKitemSideInputParDoFn.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkKitemSideInputParDoFn.java index 8dc6115f74f1..f1cae3518ac0 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkKitemSideInputParDoFn.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkKitemSideInputParDoFn.java @@ -52,7 +52,7 @@ "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) "nullness" // TODO(https://github.com/apache/beam/issues/20497) }) -/** Similar to {@link SimpleParDoFn} but for splittable ProcessFns. */ +/* Similar to {@link SimpleParDoFn} but for splittable ProcessFns. */ public class StreamingKeyedWorkKitemSideInputParDoFn implements ParDoFn { private final StateTag> keyAddr; diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputProcessor.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputProcessor.java index 477512f157f8..14fea45ef800 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputProcessor.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputProcessor.java @@ -34,10 +34,8 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; /** Helper class for handling elements blocked on side inputs. */ -@SuppressWarnings({ - "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) - "nullness" // TODO(https://github.com/apache/beam/issues/20497) -}) +@SuppressWarnings("nullness" // TODO(https://github.com/apache/beam/issues/20497) +) class StreamingSideInputProcessor { private final StreamingSideInputFetcher sideInputFetcher; diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFnHelpersTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFnHelpersTest.java index 06599870fe42..6bbbf953967d 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFnHelpersTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFnHelpersTest.java @@ -25,9 +25,14 @@ import org.apache.beam.runners.core.DoFnRunner; import org.apache.beam.runners.core.SideInputReader; import org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext; +import org.apache.beam.runners.dataflow.worker.counters.CounterFactory; +import org.apache.beam.runners.dataflow.worker.util.common.worker.Receiver; 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.windowing.GlobalWindow; +import org.apache.beam.sdk.util.DoFnInfo; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; import org.junit.Before; @@ -48,23 +53,19 @@ public class SimpleParDoFnHelpersTest { @Mock DoFnRunnerFactory runnerFactory; @Mock DoFnRunner mockRunner; - @Mock - StreamingSideInputProcessor - sideInputProcessor; + @Mock StreamingSideInputProcessor sideInputProcessor; - @Mock org.apache.beam.sdk.util.DoFnInfo doFnInfo; - @Mock org.apache.beam.runners.dataflow.worker.counters.CounterFactory counterFactory; + @Mock DoFnInfo doFnInfo; + @Mock CounterFactory counterFactory; - private static class TestDoFn extends org.apache.beam.sdk.transforms.DoFn { + private static class TestDoFn extends DoFn { @ProcessElement public void processElement() {} } private TestDoFn doFn = new TestDoFn(); - private SimpleParDoFnHelpers< - String, String, org.apache.beam.sdk.transforms.windowing.GlobalWindow> - helpers; + private SimpleParDoFnHelpers helpers; @Before @SuppressWarnings("unchecked") @@ -74,7 +75,7 @@ public void setUp() throws Exception { when(stepContext.namespacedToUser()).thenReturn(userStepContext); when(operationContext.counterFactory()).thenReturn(counterFactory); - when(doFnInstanceManager.get()).thenReturn((org.apache.beam.sdk.util.DoFnInfo) doFnInfo); + when(doFnInstanceManager.get()).thenReturn((DoFnInfo) doFnInfo); when(doFnInfo.getDoFn()).thenReturn(doFn); when(runnerFactory.createRunner( @@ -98,8 +99,7 @@ public void setUp() throws Exception { @Test public void testReallyStartBundle() throws Exception { - helpers.startBundle( - mock(org.apache.beam.runners.dataflow.worker.util.common.worker.Receiver.class)); + helpers.startBundle(mock(Receiver.class)); helpers.reallyStartBundle(); verify(runnerFactory) @@ -111,8 +111,7 @@ public void testReallyStartBundle() throws Exception { @Test public void testFinishBundle() throws Exception { - helpers.startBundle( - mock(org.apache.beam.runners.dataflow.worker.util.common.worker.Receiver.class)); + helpers.startBundle(mock(Receiver.class)); helpers.reallyStartBundle(); helpers.finishBundle(sideInputProcessor); @@ -124,8 +123,7 @@ public void testFinishBundle() throws Exception { @Test public void testAbort() throws Exception { - helpers.startBundle( - mock(org.apache.beam.runners.dataflow.worker.util.common.worker.Receiver.class)); + helpers.startBundle(mock(Receiver.class)); helpers.reallyStartBundle(); helpers.abort(); diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFnTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFnTest.java index bd742226c804..5b1d7a8d3664 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFnTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFnTest.java @@ -489,7 +489,7 @@ public void startBundle() throws Exception { } @ProcessElement - public void processElement(ProcessContext c) throws Exception { + public void processElement() throws Exception { assertThat(startCalled, equalTo(true)); assertThat(tracker.getCurrentState(), equalTo(operationContext.getProcessState())); } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkKitemSideInputParDoFnTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkKitemSideInputParDoFnTest.java index 8dc443f51443..5aefa6d14a28 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkKitemSideInputParDoFnTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkKitemSideInputParDoFnTest.java @@ -27,6 +27,7 @@ import java.util.Collections; import java.util.List; import java.util.Map; +import java.util.concurrent.Executors; import org.apache.beam.runners.core.DoFnRunner; import org.apache.beam.runners.core.InMemoryStateInternals; import org.apache.beam.runners.core.KeyedWorkItem; @@ -45,7 +46,9 @@ import org.apache.beam.runners.dataflow.worker.windmill.Windmill; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.Timer; import org.apache.beam.sdk.coders.BigEndianIntegerCoder; +import org.apache.beam.sdk.coders.ByteArrayCoder; 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.options.PipelineOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; @@ -58,6 +61,7 @@ import org.apache.beam.sdk.transforms.View; import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker; import org.apache.beam.sdk.transforms.splittabledofn.SplitResult; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.transforms.windowing.FixedWindows; import org.apache.beam.sdk.transforms.windowing.IntervalWindow; import org.apache.beam.sdk.transforms.windowing.PaneInfo; @@ -88,7 +92,7 @@ @RunWith(JUnit4.class) public class StreamingKeyedWorkKitemSideInputParDoFnTest { private static final FixedWindows WINDOW_FN = FixedWindows.of(Duration.millis(10)); - private static TupleTag> mainOutputTag = new TupleTag<>(); + private static final TupleTag> MAIN_OUTPUT_TAG = new TupleTag<>(); private final InMemoryStateInternals state = InMemoryStateInternals.forKey("a"); @@ -96,7 +100,6 @@ public class StreamingKeyedWorkKitemSideInputParDoFnTest { @Mock private TimerInternals mockTimerInternals; @Mock private SideInputReader mockSideInputReader; - @SuppressWarnings({"unchecked", "rawtypes"}) @Before public void setUp() { MockitoAnnotations.initMocks(this); @@ -111,17 +114,13 @@ public void testInvokeProcessElement() throws Exception { PCollectionView view = createView(); when(stepContext.issueSideInputFetch( - eq(view), - any(org.apache.beam.sdk.transforms.windowing.BoundedWindow.class), - eq(SideInputState.UNKNOWN))) + eq(view), any(BoundedWindow.class), eq(SideInputState.UNKNOWN))) .thenReturn(false); when(stepContext.issueSideInputFetch( - eq(view), - any(org.apache.beam.sdk.transforms.windowing.BoundedWindow.class), - eq(SideInputState.KNOWN_READY))) + eq(view), any(BoundedWindow.class), eq(SideInputState.KNOWN_READY))) .thenReturn(true); - when(mockTimerInternals.currentInputWatermarkTime()).thenReturn(new Instant(15L)); + when(mockTimerInternals.currentInputWatermarkTime()).thenReturn(Instant.ofEpochMilli(15L)); StreamingKeyedWorkKitemSideInputParDoFn, IntervalWindow> runner = createRunner(view); @@ -141,13 +140,13 @@ public void testInvokeProcessElement() throws Exception { // Initially blocked! No output. assertEquals(0, receiver.outputs.size()); - when(mockTimerInternals.currentInputWatermarkTime()).thenReturn(new Instant(20)); + when(mockTimerInternals.currentInputWatermarkTime()).thenReturn(Instant.ofEpochMilli(20)); runner.processElement( new ValueInEmptyWindows<>( KeyedWorkItems.timersWorkItem( "a", ImmutableList.of( - timerData(window(10, 20), new Instant(19), Timer.Type.WATERMARK))))); + timerData(window(10, 20), Instant.ofEpochMilli(19), Timer.Type.WATERMARK))))); // Timer is blocked too! assertEquals(0, receiver.outputs.size()); @@ -218,17 +217,13 @@ public void testSplittableProcessElement() throws Exception { PCollectionView view = createView(); when(stepContext.issueSideInputFetch( - eq(view), - any(org.apache.beam.sdk.transforms.windowing.BoundedWindow.class), - eq(SideInputState.UNKNOWN))) + eq(view), any(BoundedWindow.class), eq(SideInputState.UNKNOWN))) .thenReturn(false); when(stepContext.issueSideInputFetch( - eq(view), - any(org.apache.beam.sdk.transforms.windowing.BoundedWindow.class), - eq(SideInputState.KNOWN_READY))) + eq(view), any(BoundedWindow.class), eq(SideInputState.KNOWN_READY))) .thenReturn(true); - when(mockTimerInternals.currentInputWatermarkTime()).thenReturn(new Instant(15L)); + when(mockTimerInternals.currentInputWatermarkTime()).thenReturn(Instant.ofEpochMilli(15L)); StreamingKeyedWorkKitemSideInputParDoFn< byte[], KV, KV, IntervalWindow> runner = createSplittableRunner(view); @@ -267,7 +262,7 @@ public void testSplittableProcessElement() throws Exception { } private WindowedValue createDatum(T element, long timestampMillis) { - Instant timestamp = new Instant(timestampMillis); + Instant timestamp = Instant.ofEpochMilli(timestampMillis); return WindowedValues.of( element, timestamp, Arrays.asList(WINDOW_FN.assignWindow(timestamp)), PaneInfo.NO_FIRING); } @@ -282,7 +277,7 @@ private TimerData timerData(IntervalWindow window, Instant timestamp, Timer.Type } private IntervalWindow window(long start, long end) { - return new IntervalWindow(new Instant(start), new Instant(end)); + return new IntervalWindow(Instant.ofEpochMilli(start), Instant.ofEpochMilli(end)); } private PCollectionView createView() { @@ -328,7 +323,7 @@ public void processElement(ProcessContext c) { windowingStrategy, ImmutableList.of(view), (Coder) null, - mainOutputTag, + MAIN_OUTPUT_TAG, DoFnSchemaInformation.create(), Collections.emptyMap()); @@ -373,8 +368,8 @@ public DoFnRunner, KV> createRun options, DoFnInstanceManagers.singleInstance(fnInfo), mockSideInputReader, - mainOutputTag, - ImmutableMap.of(mainOutputTag, 0), + MAIN_OUTPUT_TAG, + ImmutableMap.of(MAIN_OUTPUT_TAG, 0), stepContext, TestOperationContext.create(), DoFnSchemaInformation.create(), @@ -388,10 +383,9 @@ public DoFnRunner, KV> createRun private StreamingKeyedWorkKitemSideInputParDoFn< byte[], KV, KV, IntervalWindow> createSplittableRunner(PCollectionView view) throws Exception { - org.apache.beam.sdk.coders.ByteArrayCoder keyCoder = - org.apache.beam.sdk.coders.ByteArrayCoder.of(); + ByteArrayCoder keyCoder = ByteArrayCoder.of(); Coder> inputCoder = - org.apache.beam.sdk.coders.KvCoder.of(BigEndianIntegerCoder.of(), StringUtf8Coder.of()); + KvCoder.of(BigEndianIntegerCoder.of(), StringUtf8Coder.of()); WindowingStrategy windowingStrategy = (WindowingStrategy) WindowingStrategy.of(WINDOW_FN); @@ -406,7 +400,7 @@ public DoFnRunner, KV> createRun (Coder) StringUtf8Coder.of(), // watermarkEstimatorStateCoder windowingStrategy, Collections.emptyMap()); - processFn.setup(org.apache.beam.sdk.options.PipelineOptionsFactory.create()); + processFn.setup(PipelineOptionsFactory.create()); DoFnInfo>, KV> fnInfo = DoFnInfo.forFn( @@ -414,7 +408,7 @@ public DoFnRunner, KV> createRun windowingStrategy, ImmutableList.of(view), (Coder) null, - mainOutputTag, + MAIN_OUTPUT_TAG, DoFnSchemaInformation.create(), Collections.emptyMap()); @@ -453,7 +447,7 @@ public DoFnRunner, KV> createRun outputManager2, mainOutputTag, sideInputReader, - java.util.concurrent.Executors.newSingleThreadScheduledExecutor(), + Executors.newSingleThreadScheduledExecutor(), 10000, Duration.standardSeconds(10), () -> null)); @@ -481,8 +475,8 @@ public DoFnRunner, KV> createRun options, DoFnInstanceManagers.singleInstance(fnInfo), mockSideInputReader, - mainOutputTag, - ImmutableMap.of(mainOutputTag, 0), + MAIN_OUTPUT_TAG, + ImmutableMap.of(MAIN_OUTPUT_TAG, 0), stepContext, TestOperationContext.create(), DoFnSchemaInformation.create(), diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputProcessorTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputProcessorTest.java index e6e553e238ce..19e22b038839 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputProcessorTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputProcessorTest.java @@ -34,10 +34,13 @@ import java.util.HashSet; import java.util.Iterator; import java.util.Set; +import org.apache.beam.runners.core.StateNamespaces; import org.apache.beam.runners.core.TimerInternals.TimerData; import org.apache.beam.sdk.state.BagState; +import org.apache.beam.sdk.state.TimeDomain; import org.apache.beam.sdk.transforms.windowing.IntervalWindow; import org.apache.beam.sdk.transforms.windowing.PaneInfo; +import org.apache.beam.sdk.values.CausedByDrain; import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.sdk.values.WindowedValues; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; @@ -79,14 +82,16 @@ public void testTryUnblockElementsNoReadyWindows() { @Test public void testTryUnblockElementsWithReadyWindows() { // Given - IntervalWindow window1 = new IntervalWindow(new Instant(0), new Instant(10)); - IntervalWindow window2 = new IntervalWindow(new Instant(10), new Instant(20)); + IntervalWindow window1 = new IntervalWindow(Instant.ofEpochMilli(0), Instant.ofEpochMilli(10)); + IntervalWindow window2 = new IntervalWindow(Instant.ofEpochMilli(10), Instant.ofEpochMilli(20)); Set readyWindows = new HashSet<>(Arrays.asList(window1, window2)); WindowedValue element1 = - WindowedValues.of("e1", new Instant(5), Arrays.asList(window1), PaneInfo.NO_FIRING); + WindowedValues.of( + "e1", Instant.ofEpochMilli(5), Arrays.asList(window1), PaneInfo.NO_FIRING); WindowedValue element2 = - WindowedValues.of("e2", new Instant(15), Arrays.asList(window2), PaneInfo.NO_FIRING); + WindowedValues.of( + "e2", Instant.ofEpochMilli(15), Arrays.asList(window2), PaneInfo.NO_FIRING); @SuppressWarnings("unchecked") BagState> mockBag1 = mock(BagState.class); @@ -131,9 +136,9 @@ public void testHandleFinishBundle() { @Test public void testHandleProcessElementBlocked() { // Given - IntervalWindow window = new IntervalWindow(new Instant(0), new Instant(10)); + IntervalWindow window = new IntervalWindow(Instant.ofEpochMilli(0), Instant.ofEpochMilli(10)); WindowedValue compressedElement = - WindowedValues.of("e", new Instant(5), Arrays.asList(window), PaneInfo.NO_FIRING); + WindowedValues.of("e", Instant.ofEpochMilli(5), Arrays.asList(window), PaneInfo.NO_FIRING); when(mockFetcher.storeIfBlocked(any(WindowedValue.class))).thenReturn(true); @@ -151,10 +156,11 @@ public void testHandleProcessElementBlocked() { @Test public void testHandleProcessElementUnblocked() { // Given - IntervalWindow window1 = new IntervalWindow(new Instant(0), new Instant(10)); - IntervalWindow window2 = new IntervalWindow(new Instant(10), new Instant(20)); + IntervalWindow window1 = new IntervalWindow(Instant.ofEpochMilli(0), Instant.ofEpochMilli(10)); + IntervalWindow window2 = new IntervalWindow(Instant.ofEpochMilli(10), Instant.ofEpochMilli(20)); WindowedValue compressedElement = - WindowedValues.of("e", new Instant(5), Arrays.asList(window1, window2), PaneInfo.NO_FIRING); + WindowedValues.of( + "e", Instant.ofEpochMilli(5), Arrays.asList(window1, window2), PaneInfo.NO_FIRING); when(mockFetcher.storeIfBlocked(any(WindowedValue.class))).thenReturn(false); @@ -174,28 +180,36 @@ public void testHandleProcessElementUnblocked() { @Test public void testHandleProcessTimerSuccess() { // Given - TimerData mockTimer = mock(TimerData.class); - when(mockFetcher.storeIfBlocked(mockTimer)).thenReturn(false); + TimerData testTimer = + TimerData.of( + StateNamespaces.global(), + Instant.ofEpochMilli(1000), + Instant.ofEpochMilli(2000), + TimeDomain.EVENT_TIME, + CausedByDrain.NORMAL); + when(mockFetcher.storeIfBlocked(testTimer)).thenReturn(false); // When - processor.handleProcessTimer(mockTimer); + processor.handleProcessTimer(testTimer); // Then - verify(mockFetcher).storeIfBlocked(mockTimer); + verify(mockFetcher).storeIfBlocked(testTimer); } @Test public void testHandleProcessTimerThrowsPreconditionFail() { // Given - TimerData mockTimer = mock(TimerData.class); - when(mockFetcher.storeIfBlocked(mockTimer)).thenReturn(true); + TimerData testTimer = + TimerData.of( + StateNamespaces.global(), + Instant.ofEpochMilli(1000), + Instant.ofEpochMilli(2000), + TimeDomain.EVENT_TIME, + CausedByDrain.NORMAL); + when(mockFetcher.storeIfBlocked(testTimer)).thenReturn(true); // When & Then - assertThrows( - IllegalStateException.class, - () -> { - processor.handleProcessTimer(mockTimer); - }); - verify(mockFetcher).storeIfBlocked(mockTimer); + assertThrows(IllegalStateException.class, () -> processor.handleProcessTimer(testTimer)); + verify(mockFetcher).storeIfBlocked(testTimer); } } 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 6bb0203231a0..43331d11a7ee 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 @@ -323,10 +323,6 @@ private CloudObject getCloudObject(DoFn fn) { private CloudObject getCloudObject(DoFn fn, WindowingStrategy windowingStrategy) { CloudObject object = CloudObject.forClassName("DoFn"); - @SuppressWarnings({ - "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) - "unchecked" - }) DoFnInfo info = DoFnInfo.forFn( fn, @@ -377,9 +373,10 @@ public void testCleanupRegistered() throws Exception { Receiver rcvr = new OutputReceiver(); parDoFn.startBundle(rcvr); - IntervalWindow firstWindow = new IntervalWindow(new Instant(0), new Instant(10)); + IntervalWindow firstWindow = + new IntervalWindow(Instant.ofEpochMilli(0), Instant.ofEpochMilli(10)); parDoFn.processElement( - WindowedValues.of("foo", new Instant(1), firstWindow, PaneInfo.NO_FIRING)); + WindowedValues.of("foo", Instant.ofEpochMilli(1), firstWindow, PaneInfo.NO_FIRING)); verify(stepContext) .setStateCleanupTimer( @@ -436,7 +433,7 @@ public void testCleanupTimerForGlobalWindowWithAllowedLateness() throws Exceptio GlobalWindow globalWindow = GlobalWindow.INSTANCE; parDoFn.processElement( - WindowedValues.of("foo", new Instant(1), globalWindow, PaneInfo.NO_FIRING)); + WindowedValues.of("foo", Instant.ofEpochMilli(1), globalWindow, PaneInfo.NO_FIRING)); assertThat( globalWindow.maxTimestamp().plus(allowedLateness), @@ -516,8 +513,10 @@ public void testCleanupWorks() throws Exception { Receiver rcvr = new OutputReceiver(); parDoFn.startBundle(rcvr); - IntervalWindow firstWindow = new IntervalWindow(new Instant(0), new Instant(9)); - IntervalWindow secondWindow = new IntervalWindow(new Instant(10), new Instant(19)); + IntervalWindow firstWindow = + new IntervalWindow(Instant.ofEpochMilli(0), Instant.ofEpochMilli(9)); + IntervalWindow secondWindow = + new IntervalWindow(Instant.ofEpochMilli(10), Instant.ofEpochMilli(19)); Coder windowCoder = IntervalWindow.getCoder(); StateNamespace firstWindowNamespace = StateNamespaces.window(windowCoder, firstWindow); From 7ad6c7b98b4e77a1e0ea90733ac26503c458c5f6 Mon Sep 17 00:00:00 2001 From: Reuven Lax Date: Thu, 4 Jun 2026 10:13:02 -0700 Subject: [PATCH 09/11] fix windows --- .../dataflow/worker/SimpleParDoFn.java | 4 +- .../dataflow/worker/SimpleParDoFnHelpers.java | 15 +++-- ...reamingKeyedWorkItemSideInputParDoFn.java} | 66 +++++++++---------- .../worker/StreamingSideInputProcessor.java | 5 +- .../dataflow/worker/UserParDoFnFactory.java | 2 +- ...ingKeyedWorkItemSideInputParDoFnTest.java} | 16 ++--- 6 files changed, 54 insertions(+), 54 deletions(-) rename runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/{StreamingKeyedWorkKitemSideInputParDoFn.java => StreamingKeyedWorkItemSideInputParDoFn.java} (82%) rename runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/{StreamingKeyedWorkKitemSideInputParDoFnTest.java => StreamingKeyedWorkItemSideInputParDoFnTest.java} (97%) 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 90492d4ee8d3..34dff6b88358 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 @@ -186,13 +186,13 @@ public void processTimers() throws Exception { helpers.userStepContext, windowCoder, this::onStartKey, - sideInputProcessor); + () -> sideInputProcessor); helpers.processTimers( SimpleParDoFnHelpers.TimerType.SYSTEM, helpers.stepContext, windowCoder, this::onStartKey, - sideInputProcessor); + () -> sideInputProcessor); } @Override 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 df0af65a4b9e..8398fe8596d0 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 @@ -25,6 +25,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.function.Supplier; import org.apache.beam.runners.core.DoFnRunner; import org.apache.beam.runners.core.SideInputReader; import org.apache.beam.runners.core.StateInternals; @@ -260,9 +261,9 @@ enum TimerType { public void processTimer( SimpleParDoFnHelpers doFn, TimerInternals.TimerData timer, - StreamingSideInputProcessor sideInputProcessor) + Supplier sideInputProcessor) throws Exception { - doFn.processUserTimer(timer, sideInputProcessor); + doFn.processUserTimer(timer, sideInputProcessor.get()); } }, FAIL_USER { @@ -270,7 +271,7 @@ public void processTimer( public void processTimer( SimpleParDoFnHelpers doFn, TimerInternals.TimerData timer, - StreamingSideInputProcessor sideInputProcessor) + Supplier sideInputProcessor) throws Exception { throw new UnsupportedOperationException( "Attempt to deliver a timer to a DoFn, but timers are not supported here."); @@ -281,16 +282,16 @@ public void processTimer( public void processTimer( SimpleParDoFnHelpers doFn, TimerInternals.TimerData timer, - StreamingSideInputProcessor sideInputProcessor) + Supplier sideInputProcessor) throws Exception { - doFn.processSystemTimer(timer, sideInputProcessor); + doFn.processSystemTimer(timer, sideInputProcessor.get()); } }; public abstract void processTimer( SimpleParDoFnHelpers doFn, TimerInternals.TimerData timer, - StreamingSideInputProcessor sideInputProcessor) + Supplier sideInputProcessor) throws Exception; }; @@ -299,7 +300,7 @@ void processTimers( DataflowExecutionContext.DataflowStepContext context, Coder windowCoder, Runnable startKey, - StreamingSideInputProcessor sideInputProcessor) + Supplier> sideInputProcessor) throws Exception { TimerInternals.TimerData timer = context.getNextFiredTimer(windowCoder); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkKitemSideInputParDoFn.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputParDoFn.java similarity index 82% rename from runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkKitemSideInputParDoFn.java rename to runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputParDoFn.java index f1cae3518ac0..225bc6af0ea9 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkKitemSideInputParDoFn.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputParDoFn.java @@ -53,14 +53,14 @@ "nullness" // TODO(https://github.com/apache/beam/issues/20497) }) /* Similar to {@link SimpleParDoFn} but for splittable ProcessFns. */ -public class StreamingKeyedWorkKitemSideInputParDoFn +public class StreamingKeyedWorkItemSideInputParDoFn implements ParDoFn { private final StateTag> keyAddr; private final Coder inputCoder; private final SimpleParDoFnHelpers, OutputT, W> helpers; protected @Nullable StreamingSideInputProcessor sideInputProcessor; - StreamingKeyedWorkKitemSideInputParDoFn( + StreamingKeyedWorkItemSideInputParDoFn( PipelineOptions options, DoFnInstanceManager doFnInstanceManager, SideInputReader sideInputReader, @@ -120,27 +120,28 @@ protected void onStartKey() { boolean hasState = helpers.hasState(); // TODO(relax): We should be able to get this without writing it to state! - K key = keyValue().read(); - - sideInputProcessor.tryUnblockElementsAndTimers( - (unblockedElements, unblockedTimers) -> { - if (!Iterables.isEmpty(unblockedElements) || !Iterables.isEmpty(unblockedTimers)) { - helpers.fnRunner.processElement( - new ValueInEmptyWindows<>( - KeyedWorkItems.workItem(key, unblockedTimers, unblockedElements))); - } - if (hasState) { - List windows = - (List) - StreamSupport.stream(unblockedElements.spliterator(), false) - .flatMap(wv -> wv.getWindows().stream()) - .collect(Collectors.toList()); - // These elements are now processed. Register cleanup timers for all the unblocked - // windows. - helpers.registerStateCleanup( - (WindowingStrategy) getDoFnInfo().getWindowingStrategy(), windows); - } - }); + @Nullable K key = keyValue().read(); + if (key != null) { + sideInputProcessor.tryUnblockElementsAndTimers( + (unblockedElements, unblockedTimers) -> { + if (!Iterables.isEmpty(unblockedElements) || !Iterables.isEmpty(unblockedTimers)) { + helpers.fnRunner.processElement( + new ValueInEmptyWindows<>( + KeyedWorkItems.workItem(key, unblockedTimers, unblockedElements))); + } + if (hasState) { + List windows = + (List) + StreamSupport.stream(unblockedElements.spliterator(), false) + .flatMap(wv -> wv.getWindows().stream()) + .collect(Collectors.toList()); + // These elements are now processed. Register cleanup timers for all the unblocked + // windows. + helpers.registerStateCleanup( + (WindowingStrategy) getDoFnInfo().getWindowingStrategy(), windows); + } + }); + } } } @@ -183,13 +184,13 @@ public void processTimers() throws Exception { helpers.userStepContext, windowCoder, this::onStartKey, - sideInputProcessor); + () -> sideInputProcessor); helpers.processTimers( SimpleParDoFnHelpers.TimerType.SYSTEM, helpers.stepContext, windowCoder, this::onStartKey, - sideInputProcessor); + () -> sideInputProcessor); } @Override @@ -212,17 +213,14 @@ protected void onProcessWindowedValue(WindowedValue> el Collection windowsProcessed; if (sideInputProcessor != null) { windowsProcessed = hasState ? Lists.newArrayList() : Collections.emptyList(); - KeyedWorkItem unblocked = sideInputProcessor.handleProcessKeyedWorkItem(elem); - if (!Iterables.isEmpty(unblocked.elementsIterable()) - || !Iterables.isEmpty(unblocked.timersIterable())) { - helpers.fnRunner.processElement(elem.withValue(unblocked)); + WindowedValue> unblocked = + sideInputProcessor.handleProcessKeyedWorkItem(elem); + if (!Iterables.isEmpty(unblocked.getValue().elementsIterable()) + || !Iterables.isEmpty(unblocked.getValue().timersIterable())) { + helpers.fnRunner.processElement(unblocked); } if (hasState) { - windowsProcessed = - (Collection) - StreamSupport.stream(unblocked.elementsIterable().spliterator(), false) - .flatMap(wv -> wv.getWindows().stream()) - .collect(Collectors.toList()); + windowsProcessed.addAll((Collection) unblocked.getWindows()); } } else { helpers.fnRunner.processElement(elem); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputProcessor.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputProcessor.java index 14fea45ef800..34c1a06d54de 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputProcessor.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputProcessor.java @@ -102,7 +102,7 @@ Iterator> handleProcessElement( (WindowedValue e) -> !sideInputFetcher.storeIfBlocked(e)); } - KeyedWorkItem handleProcessKeyedWorkItem( + WindowedValue> handleProcessKeyedWorkItem( WindowedValue> elem) { List> readyInputs = Lists.newArrayList( @@ -117,7 +117,8 @@ KeyedWorkItem handleProcessKeyedWorkItem( timer -> !sideInputFetcher.storeIfBlocked(timer))); KeyedWorkItem keyedWorkItem = KeyedWorkItems.workItem(elem.getValue().key(), readyTimers, readyInputs); - return keyedWorkItem; + + return elem.withValue(keyedWorkItem); } void handleProcessTimer(TimerInternals.TimerData timer) { diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UserParDoFnFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UserParDoFnFactory.java index 6807a2ecbd5c..9466ad60d414 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UserParDoFnFactory.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UserParDoFnFactory.java @@ -159,7 +159,7 @@ public ParDoFn create( if (streamingKeyedWorkItem && hasStreamingSideInput) { KeyedWorkItemCoder> kwiCoder = (KeyedWorkItemCoder>) doFnInfo.getInputCoder(); - return new StreamingKeyedWorkKitemSideInputParDoFn<>( + return new StreamingKeyedWorkItemSideInputParDoFn<>( options, instanceManager, sideInputReader, diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkKitemSideInputParDoFnTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputParDoFnTest.java similarity index 97% rename from runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkKitemSideInputParDoFnTest.java rename to runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputParDoFnTest.java index 5aefa6d14a28..2fed6fd405a5 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkKitemSideInputParDoFnTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputParDoFnTest.java @@ -88,9 +88,9 @@ import org.mockito.Mock; import org.mockito.MockitoAnnotations; -/** Unit tests for {@link StreamingKeyedWorkKitemSideInputParDoFn}. */ +/** Unit tests for {@link StreamingKeyedWorkItemSideInputParDoFn}. */ @RunWith(JUnit4.class) -public class StreamingKeyedWorkKitemSideInputParDoFnTest { +public class StreamingKeyedWorkItemSideInputParDoFnTest { private static final FixedWindows WINDOW_FN = FixedWindows.of(Duration.millis(10)); private static final TupleTag> MAIN_OUTPUT_TAG = new TupleTag<>(); @@ -121,7 +121,7 @@ public void testInvokeProcessElement() throws Exception { .thenReturn(true); when(mockTimerInternals.currentInputWatermarkTime()).thenReturn(Instant.ofEpochMilli(15L)); - StreamingKeyedWorkKitemSideInputParDoFn, IntervalWindow> + StreamingKeyedWorkItemSideInputParDoFn, IntervalWindow> runner = createRunner(view); TestReceiver receiver = new TestReceiver(); @@ -224,7 +224,7 @@ public void testSplittableProcessElement() throws Exception { .thenReturn(true); when(mockTimerInternals.currentInputWatermarkTime()).thenReturn(Instant.ofEpochMilli(15L)); - StreamingKeyedWorkKitemSideInputParDoFn< + StreamingKeyedWorkItemSideInputParDoFn< byte[], KV, KV, IntervalWindow> runner = createSplittableRunner(view); @@ -298,7 +298,7 @@ public void process(Object outputElem) { } @SuppressWarnings("unchecked") - private StreamingKeyedWorkKitemSideInputParDoFn< + private StreamingKeyedWorkItemSideInputParDoFn< String, Integer, KV, IntervalWindow> createRunner(PCollectionView view) throws Exception { Coder keyCoder = StringUtf8Coder.of(); @@ -364,7 +364,7 @@ public DoFnRunner, KV> createRun PipelineOptions options = PipelineOptionsFactory.create(); options.as(StreamingOptions.class).setStreaming(true); - return new StreamingKeyedWorkKitemSideInputParDoFn<>( + return new StreamingKeyedWorkItemSideInputParDoFn<>( options, DoFnInstanceManagers.singleInstance(fnInfo), mockSideInputReader, @@ -380,7 +380,7 @@ public DoFnRunner, KV> createRun } @SuppressWarnings("unchecked") - private StreamingKeyedWorkKitemSideInputParDoFn< + private StreamingKeyedWorkItemSideInputParDoFn< byte[], KV, KV, IntervalWindow> createSplittableRunner(PCollectionView view) throws Exception { ByteArrayCoder keyCoder = ByteArrayCoder.of(); @@ -471,7 +471,7 @@ public DoFnRunner, KV> createRun PipelineOptions options = PipelineOptionsFactory.create(); options.as(StreamingOptions.class).setStreaming(true); - return new StreamingKeyedWorkKitemSideInputParDoFn<>( + return new StreamingKeyedWorkItemSideInputParDoFn<>( options, DoFnInstanceManagers.singleInstance(fnInfo), mockSideInputReader, From 662a000c367596795952a97bbbce39a820015d79 Mon Sep 17 00:00:00 2001 From: Reuven Lax Date: Thu, 4 Jun 2026 10:33:52 -0700 Subject: [PATCH 10/11] fix compilation --- .../runners/dataflow/worker/SimpleParDoFnHelpers.java | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) 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 8398fe8596d0..964cf2323d51 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 @@ -261,7 +261,7 @@ enum TimerType { public void processTimer( SimpleParDoFnHelpers doFn, TimerInternals.TimerData timer, - Supplier sideInputProcessor) + Supplier> sideInputProcessor) throws Exception { doFn.processUserTimer(timer, sideInputProcessor.get()); } @@ -271,7 +271,7 @@ public void processTimer( public void processTimer( SimpleParDoFnHelpers doFn, TimerInternals.TimerData timer, - Supplier sideInputProcessor) + Supplier> sideInputProcessor) throws Exception { throw new UnsupportedOperationException( "Attempt to deliver a timer to a DoFn, but timers are not supported here."); @@ -282,7 +282,7 @@ public void processTimer( public void processTimer( SimpleParDoFnHelpers doFn, TimerInternals.TimerData timer, - Supplier sideInputProcessor) + Supplier> sideInputProcessor) throws Exception { doFn.processSystemTimer(timer, sideInputProcessor.get()); } @@ -291,7 +291,7 @@ public void processTimer( public abstract void processTimer( SimpleParDoFnHelpers doFn, TimerInternals.TimerData timer, - Supplier sideInputProcessor) + Supplier> sideInputProcessor) throws Exception; }; From c9023874c7efc9948652edc6bab7d7e21d461e19 Mon Sep 17 00:00:00 2001 From: Reuven Lax Date: Thu, 4 Jun 2026 12:51:02 -0700 Subject: [PATCH 11/11] foo --- .../java/org/apache/beam/sdk/transforms/SplittableDoFnTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/SplittableDoFnTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/SplittableDoFnTest.java index 117ceac32d71..80d8728aa01b 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/SplittableDoFnTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/SplittableDoFnTest.java @@ -426,7 +426,7 @@ private static SDFWithSideInputBase sdfWithSideInput( @Test @Category({ValidatesRunner.class, UsesBoundedSplittableParDo.class, UsesSideInputs.class}) - public void tBounded() { + public void testSideInputBounded() { testSideInput(IsBounded.BOUNDED); }