diff --git a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark4.json b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark4.json new file mode 100644 index 000000000000..e3d6056a5de9 --- /dev/null +++ b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark4.json @@ -0,0 +1,4 @@ +{ + "comment": "Modify this file in a trivial way to cause this test suite to run", + "modification": 1 +} diff --git a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json index 77f63217b86d..cad8d98b8ea5 100644 --- a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json +++ b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json @@ -7,5 +7,6 @@ "https://github.com/apache/beam/pull/34123": "noting that PR #34123 should run this test", "https://github.com/apache/beam/pull/34080": "noting that PR #34080 should run this test", "https://github.com/apache/beam/pull/34155": "noting that PR #34155 should run this test", - "https://github.com/apache/beam/pull/35159": "moving WindowedValue and making an interface" + "https://github.com/apache/beam/pull/35159": "moving WindowedValue and making an interface", + "https://github.com/apache/beam/pull/39793": "noting that PR #39793 should run this test" } diff --git a/CHANGES.md b/CHANGES.md index 73966a48313c..56c062790e22 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -69,6 +69,7 @@ ## New Features / Improvements * X feature added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). +* (Java) Spark Structured Streaming runner: stateful ParDo with state, timers, `@RequiresTimeSortedInput` and tagged outputs is now supported in batch mode ([#39779](https://github.com/apache/beam/issues/39779)). ## Breaking Changes diff --git a/runners/spark/spark_runner.gradle b/runners/spark/spark_runner.gradle index 77da3d36db92..2a161db0a820 100644 --- a/runners/spark/spark_runner.gradle +++ b/runners/spark/spark_runner.gradle @@ -510,15 +510,9 @@ tasks.register("validatesStructuredStreamingRunnerBatch", Test) { excludeCategories 'org.apache.beam.sdk.testing.UsesUnboundedPCollections' excludeCategories 'org.apache.beam.sdk.testing.UsesTestStream' // State and Timers - excludeCategories 'org.apache.beam.sdk.testing.UsesStatefulParDo' - excludeCategories 'org.apache.beam.sdk.testing.UsesMapState' - excludeCategories 'org.apache.beam.sdk.testing.UsesMultimapState' - excludeCategories 'org.apache.beam.sdk.testing.UsesSetState' - excludeCategories 'org.apache.beam.sdk.testing.UsesOrderedListState' - excludeCategories 'org.apache.beam.sdk.testing.UsesTimersInParDo' - excludeCategories 'org.apache.beam.sdk.testing.UsesTimerMap' - excludeCategories 'org.apache.beam.sdk.testing.UsesKeyInParDo' excludeCategories 'org.apache.beam.sdk.testing.UsesOnWindowExpiration' + // Every UsesOrderedListState test also uses @OnWindowExpiration, which is unsupported + excludeCategories 'org.apache.beam.sdk.testing.UsesOrderedListState' // Metrics excludeCategories 'org.apache.beam.sdk.testing.UsesCommittedMetrics' excludeCategories 'org.apache.beam.sdk.testing.UsesSystemMetrics' @@ -532,6 +526,11 @@ tasks.register("validatesStructuredStreamingRunnerBatch", Test) { excludeCategories 'org.apache.beam.sdk.testing.UsesTriggeredSideInputs' } filter { + // These build on PeriodicImpulse, so the pipeline is unbounded and rejected by this batch only + // runner, but they are not categorized as UsesUnboundedPCollections. Excluded by name rather + // than adding that category upstream, which would also stop other runners running them. + excludeTestsMatching 'org.apache.beam.sdk.transforms.PerKeyOrderingTest.testMultipleStatefulOrderingWithShuffle' + excludeTestsMatching 'org.apache.beam.sdk.transforms.PerKeyOrderingTest.testMultipleStatefulOrderingWithoutShuffle' // Combine with context not implemented excludeTestsMatching 'org.apache.beam.sdk.transforms.CombineFnsTest.testComposedCombineWithContext' excludeTestsMatching 'org.apache.beam.sdk.transforms.CombineTest$CombineWithContextTests.testSimpleCombineWithContext' diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/DoFnRunnerFactory.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/DoFnRunnerFactory.java index 5e8703a05b06..ce4155ee8e19 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/DoFnRunnerFactory.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/DoFnRunnerFactory.java @@ -25,6 +25,7 @@ import org.apache.beam.runners.core.DoFnRunner; import org.apache.beam.runners.core.DoFnRunners; import org.apache.beam.runners.core.SideInputReader; +import org.apache.beam.runners.core.StepContext; import org.apache.beam.runners.spark.structuredstreaming.metrics.MetricsAccumulator; import org.apache.beam.runners.spark.structuredstreaming.translation.batch.functions.CachedSideInputReader; import org.apache.beam.runners.spark.structuredstreaming.translation.batch.functions.NoOpStepContext; @@ -71,6 +72,20 @@ interface DoFnRunnerWithTeardown extends DoFnRunner { abstract DoFnRunnerWithTeardown create( PipelineOptions options, MetricsAccumulator metrics, WindowedValueMultiReceiver output); + /** + * Creates a runner backed by {@code stepContext} so that state and timers are available. + * + *

Only supported for a single, unfused {@link DoFn}: a fused runner cannot drive timers. + */ + DoFnRunnerWithTeardown create( + PipelineOptions options, + MetricsAccumulator metrics, + WindowedValueMultiReceiver output, + StepContext stepContext) { + throw new UnsupportedOperationException( + "Stateful execution is not supported by " + getClass().getSimpleName()); + } + /** * Fuses the factory for the following {@link DoFnRunner} into a single factory that processes * both DoFns in a single step. @@ -128,6 +143,15 @@ DoFnRunnerFactory fuse(DoFnRunnerFactory next) { @Override DoFnRunnerWithTeardown create( PipelineOptions options, MetricsAccumulator metrics, WindowedValueMultiReceiver output) { + return create(options, metrics, output, new NoOpStepContext()); + } + + @Override + DoFnRunnerWithTeardown create( + PipelineOptions options, + MetricsAccumulator metrics, + WindowedValueMultiReceiver output, + StepContext stepContext) { DoFnRunner simpleRunner = DoFnRunners.simpleRunner( options, @@ -136,7 +160,7 @@ DoFnRunnerWithTeardown create( filterMainOutput ? new FilteredOutput<>(output, mainOutput) : output, mainOutput, additionalOutputs, - new NoOpStepContext(), + stepContext, coder, outputCoders, windowingStrategy, diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/ParDoTranslatorBatch.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/ParDoTranslatorBatch.java index 0f43f329b0df..cb05da63dc3e 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/ParDoTranslatorBatch.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/ParDoTranslatorBatch.java @@ -35,6 +35,7 @@ import org.apache.beam.runners.core.SideInputReader; import org.apache.beam.runners.spark.SparkCommonPipelineOptions; import org.apache.beam.runners.spark.structuredstreaming.metrics.MetricsAccumulator; +import org.apache.beam.runners.spark.structuredstreaming.translation.PipelineTranslator.TranslationState; import org.apache.beam.runners.spark.structuredstreaming.translation.PipelineTranslator.UnresolvedTranslation; import org.apache.beam.runners.spark.structuredstreaming.translation.TransformTranslator; import org.apache.beam.runners.spark.structuredstreaming.translation.batch.functions.SideInputValues; @@ -64,8 +65,11 @@ * *

Each tag is encoded as individual column with a respective schema & encoder each. * + *

Stateful {@link org.apache.beam.sdk.transforms.DoFn DoFns}, those using timers, and those + * annotated with {@link DoFn.RequiresTimeSortedInput} are translated by {@link + * StatefulParDoTranslatorBatch} instead. + * *

TODO: - *

  • Add support for state and timers. *
  • Add support for SplittableDoFn */ class ParDoTranslatorBatch @@ -87,18 +91,18 @@ public boolean canTranslate(ParDo.MultiOutput transform) { "Not expected to directly translate splittable DoFn, should have been overridden: %s", doFn); - // TODO: add support of states and timers + // Stateful, timer using and time sorted DoFns are routed to StatefulParDoTranslatorBatch by + // PipelineTranslatorBatch#getTransformTranslator. Reaching here with one means dispatch is + // broken, not that the feature is unsupported. checkState( - !signature.usesState() && !signature.usesTimers(), - "States and timers are not supported for the moment."); + !StatefulParDoTranslatorBatch.appliesTo(transform), + "Stateful / time sorted DoFn should have been translated by %s: %s", + StatefulParDoTranslatorBatch.class.getSimpleName(), + doFn); checkState( signature.onWindowExpiration() == null, "onWindowExpiration is not supported: %s", doFn); - checkState( - !signature.processElement().requiresTimeSortedInput(), - "@RequiresTimeSortedInput is not supported for the moment"); - SparkSideInputReader.validateMaterializations(transform.getSideInputs().values()); return true; } @@ -211,11 +215,11 @@ public Dataset> resolve( *

    This can help to avoid unnecessary caching in case of multiple outputs if only {@code * mainTag} is consumed. */ - private Map, PCollection> skipUnconsumedOutputs( + static Map, PCollection> skipUnconsumedOutputs( Map, PCollection> outputs, TupleTag mainTag, TupleTagList otherTags, - Context cxt) { + TranslationState cxt) { switch (outputs.size()) { case 1: return outputs; // always keep main output @@ -235,7 +239,7 @@ private Map, PCollection> skipUnconsumedOutputs( } } - private Map tagsColumnIndex(Collection> tags) { + static Map tagsColumnIndex(Collection> tags) { Map index = Maps.newHashMapWithExpectedSize(tags.size()); for (TupleTag tag : tags) { index.put(tag.getId(), index.size()); diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/PipelineTranslatorBatch.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/PipelineTranslatorBatch.java index c4a18801ccba..d6d85d977c04 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/PipelineTranslatorBatch.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/PipelineTranslatorBatch.java @@ -82,11 +82,27 @@ public class PipelineTranslatorBatch extends PipelineTranslator { SplittableParDo.PrimitiveBoundedRead.class, new ReadSourceTranslatorBatch<>()); } + /** + * Translators that shadow the {@link #TRANSFORM_TRANSLATORS} entry for their transform class when + * a predicate matches, so that a single transform class can be translated in more than one way + * depending on the transform instance. + * + *

    Currently only {@link ParDo.MultiOutput} needs this, to route stateful and time sorted + * {@link org.apache.beam.sdk.transforms.DoFn DoFns} away from {@link ParDoTranslatorBatch}. + */ + @SuppressWarnings("rawtypes") + private static final TransformTranslator STATEFUL_PARDO_TRANSLATOR = + new StatefulParDoTranslatorBatch<>(); + /** Returns a {@link TransformTranslator} for the given {@link PTransform} if known. */ @Override @Nullable protected > TransformTranslator getTransformTranslator(TransformT transform) { + if (transform instanceof ParDo.MultiOutput + && StatefulParDoTranslatorBatch.appliesTo((ParDo.MultiOutput) transform)) { + return STATEFUL_PARDO_TRANSLATOR; + } return TRANSFORM_TRANSLATORS.get(transform.getClass()); } } diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/StatefulDoFnGroupFunction.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/StatefulDoFnGroupFunction.java new file mode 100644 index 000000000000..77eb0a7cb65d --- /dev/null +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/StatefulDoFnGroupFunction.java @@ -0,0 +1,392 @@ +/* + * 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.spark.structuredstreaming.translation.batch; + +import static org.apache.beam.runners.spark.structuredstreaming.translation.utils.ScalaInterop.tuple; + +import java.util.ArrayDeque; +import java.util.Deque; +import java.util.Iterator; +import java.util.Map; +import java.util.function.Supplier; +import javax.annotation.CheckForNull; +import org.apache.beam.runners.core.InMemoryStateInternals; +import org.apache.beam.runners.core.InMemoryTimerInternals; +import org.apache.beam.runners.core.StateInternals; +import org.apache.beam.runners.core.StateNamespaces; +import org.apache.beam.runners.core.StepContext; +import org.apache.beam.runners.core.TimerInternals; +import org.apache.beam.runners.core.TimerInternals.TimerData; +import org.apache.beam.runners.spark.structuredstreaming.metrics.MetricsAccumulator; +import org.apache.beam.runners.spark.structuredstreaming.translation.batch.DoFnRunnerFactory.DoFnRunnerWithTeardown; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +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.TupleTag; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.AbstractIterator; +import org.apache.spark.TaskContext; +import org.apache.spark.api.java.function.FlatMapGroupsFunction; +import org.apache.spark.util.TaskCompletionListener; +import org.checkerframework.checker.nullness.qual.NonNull; +import org.checkerframework.checker.nullness.qual.Nullable; +import scala.Tuple2; + +/** + * Runs a stateful {@link DoFn} over the key groups of a {@code flatMapSortedGroups}, where the + * elements of each group are already ordered by event time. + * + *

    State is a plain heap object scoped to the key being processed and is dropped when that key is + * done. This is what makes batch state cheap: there is no state store to bridge onto, so every Beam + * state type is already implemented by {@link InMemoryStateInternals}. Memory is bounded by the + * state a single key holds, not by the number of elements that key received. + * + *

    The {@link DoFn} is set up once per task and torn down from a task completion listener, + * not once per key. Spark calls {@link #call} per key group, but {@code @Setup}/{@code @Teardown} + * bracket the lifetime of the {@link DoFn} instance, and there is a single instance per + * deserialized closure, so tearing it down between keys would violate the contract that no method + * runs after {@code @Teardown} (and would re-run expensive setup for every key). Each key gets its + * own bundle, which is the level the model does allow to vary, and its own state and + * timers via {@link MutableStepContext}. + * + *

    Outputs are pulled lazily: the {@link DoFn} pushes into a buffer and the returned iterator + * drains it, advancing the input only when the buffer runs dry, so neither a key with many elements + * nor one with many timers is ever materialized. + */ +abstract class StatefulDoFnGroupFunction, OutT extends @NonNull Object> + implements FlatMapGroupsFunction, OutT> { + + private final Supplier options; + private final MetricsAccumulator metrics; + private final DoFnRunnerFactory factory; + + private transient @Nullable Deque buffer; + private transient @Nullable MutableStepContext stepContext; + private transient @Nullable DoFnRunnerWithTeardown doFnRunner; + private transient boolean needsBundleStart; + private transient boolean isTornDown; + + private StatefulDoFnGroupFunction( + Supplier options, + MetricsAccumulator metrics, + DoFnRunnerFactory factory) { + this.options = options; + this.metrics = metrics; + this.factory = factory; + } + + /** + * {@link StatefulDoFnGroupFunction} emitting a single output of type {@link WindowedValue} of + * {@link FnOutT}. + */ + static , FnOutT> + StatefulDoFnGroupFunction> singleOutput( + Supplier options, + MetricsAccumulator metrics, + DoFnRunnerFactory factory) { + return new SingleOut<>(options, metrics, factory); + } + + /** + * {@link StatefulDoFnGroupFunction} emitting multiple outputs encoded as tuple of column index + * and {@link WindowedValue} of {@link OutT}, where column index corresponds to the index of a + * {@link TupleTag#getId()} in {@code tagColIdx}. + */ + static , FnOutT, OutT> + StatefulDoFnGroupFunction>> multiOutput( + Supplier options, + MetricsAccumulator metrics, + DoFnRunnerFactory factory, + Map tagColIdx) { + return new MultiOut<>(options, metrics, factory, tagColIdx); + } + + @Override + public Iterator call(K key, Iterator> values) { + DoFnRunnerWithTeardown runner = runner(); + // Fresh state and timers for this key; the DoFn instance itself is untouched. + stepContext().reset(key); + if (needsBundleStart) { + needsBundleStart = false; + runner.startBundle(); + } + return new StatefulGroupIt(key, values, runner); + } + + /** + * The runner for this task, created on first use. {@code factory.create} invokes {@code @Setup} + * and opens the first bundle, so this happens exactly once per task rather than once per key. + */ + private DoFnRunnerWithTeardown runner() { + DoFnRunnerWithTeardown runner = doFnRunner; + if (runner == null) { + MutableStepContext ctx = new MutableStepContext(); + Deque buf = new ArrayDeque<>(); + buffer = buf; + stepContext = ctx; + runner = factory.create(options.get(), metrics, outputManager(buf), ctx); + doFnRunner = runner; + // Spark is free to abandon an iterator part way through (a downstream limit, a task kill, an + // exception elsewhere in the stage). Tearing down from the task completion listener is the + // only way to guarantee @Teardown runs and DoFn resources are released. + TaskContext taskContext = TaskContext.get(); + if (taskContext != null) { + // An explicit listener rather than a lambda: TaskContext overloads this for both the Scala + // function and the Java interface, so a lambda is ambiguous. + taskContext.addTaskCompletionListener( + new TaskCompletionListener() { + @Override + public void onTaskCompletion(TaskContext context) { + teardownOnce(); + } + }); + } + } + return runner; + } + + private MutableStepContext stepContext() { + MutableStepContext ctx = stepContext; + if (ctx == null) { + throw new IllegalStateException("StepContext requested before the runner was created"); + } + return ctx; + } + + private Deque buffer() { + Deque buf = buffer; + if (buf == null) { + throw new IllegalStateException("Buffer requested before the runner was created"); + } + return buf; + } + + private void teardownOnce() { + DoFnRunnerWithTeardown runner = doFnRunner; + if (runner != null && !isTornDown) { + isTornDown = true; + runner.teardown(); + } + } + + /** Output manager emitting outputs of type {@link OutT} to the buffer. */ + abstract WindowedValueMultiReceiver outputManager(Deque buffer); + + /** + * {@link StatefulDoFnGroupFunction} emitting a single output of type {@link WindowedValue} of + * {@link FnOutT}. + */ + private static class SingleOut, FnOutT> + extends StatefulDoFnGroupFunction> { + private SingleOut( + Supplier options, + MetricsAccumulator metrics, + DoFnRunnerFactory factory) { + super(options, metrics, factory); + } + + @Override + WindowedValueMultiReceiver outputManager(Deque> buffer) { + return new WindowedValueMultiReceiver() { + @Override + public void output(TupleTag tag, WindowedValue output) { + buffer.add((WindowedValue) output); + } + }; + } + } + + /** + * {@link StatefulDoFnGroupFunction} emitting multiple outputs encoded as tuple of column index + * and {@link WindowedValue} of {@link OutT}, where column index corresponds to the index of a + * {@link TupleTag#getId()} in {@link #tagColIdx}. + */ + private static class MultiOut, FnOutT, OutT> + extends StatefulDoFnGroupFunction>> { + private final Map tagColIdx; + + private MultiOut( + Supplier options, + MetricsAccumulator metrics, + DoFnRunnerFactory factory, + Map tagColIdx) { + super(options, metrics, factory); + this.tagColIdx = tagColIdx; + } + + @Override + WindowedValueMultiReceiver outputManager(Deque>> buffer) { + return new WindowedValueMultiReceiver() { + @Override + public void output(TupleTag tag, WindowedValue output) { + // Additional unused outputs can be skipped here. In that case columnIdx is null. + Integer columnIdx = tagColIdx.get(tag.getId()); + if (columnIdx != null) { + buffer.add(tuple(columnIdx, (WindowedValue) output)); + } + } + }; + } + } + + /** + * A {@link StepContext} whose state and timers are swapped per key, so that one {@link DoFn} and + * one {@link org.apache.beam.runners.core.DoFnRunner DoFnRunner} can serve every key of a task. + * + *

    {@code SimpleDoFnRunner} re-reads {@code stateInternals()} on each access rather than + * caching it, which is what makes rebinding safe. + */ + private static class MutableStepContext implements StepContext { + private @Nullable StateInternals stateInternals; + private @Nullable InMemoryTimerInternals timerInternals; + + void reset(@Nullable Object key) { + stateInternals = InMemoryStateInternals.forKey(key); + timerInternals = new InMemoryTimerInternals(); + } + + InMemoryTimerInternals timers() { + InMemoryTimerInternals timers = timerInternals; + if (timers == null) { + throw new IllegalStateException("StepContext used before reset"); + } + return timers; + } + + @Override + public StateInternals stateInternals() { + StateInternals state = stateInternals; + if (state == null) { + throw new IllegalStateException("StepContext used before reset"); + } + return state; + } + + @Override + public TimerInternals timerInternals() { + return timers(); + } + } + + private class StatefulGroupIt extends AbstractIterator { + private final Iterator> groupIt; + private final K key; + private final DoFnRunnerWithTeardown runner; + private final InMemoryTimerInternals timerInternals; + + private boolean areTimersDrained; + private boolean clocksAdvanced; + private boolean isBundleFinished; + + private StatefulGroupIt( + K key, Iterator> groupIt, DoFnRunnerWithTeardown runner) { + this.key = key; + this.groupIt = groupIt; + this.runner = runner; + this.timerInternals = stepContext().timers(); + } + + @Override + protected @CheckForNull OutT computeNext() { + Deque buffer = buffer(); + try { + while (true) { + if (!buffer.isEmpty()) { + return buffer.remove(); + } + if (groupIt.hasNext()) { + runner.processElement(groupIt.next()); + } else if (!areTimersDrained) { + // Timers fire while the bundle is still open (the model processes a key's timers + // before finishBundle) and one at a time, so their output is pulled lazily too. + areTimersDrained = !fireNextTimer(); + } else if (!isBundleFinished) { + isBundleFinished = true; + needsBundleStart = true; // the next key opens a fresh bundle + runner.finishBundle(); // may produce more output + } else { + return endOfData(); // teardown is task scoped, not per key + } + } + } catch (RuntimeException re) { + teardownOnce(); + throw re; + } catch (Exception e) { + teardownOnce(); + throw new RuntimeException(e); + } + } + + /** + * Fires at most one pending timer, returning whether one fired. + * + *

    Polled once per {@code computeNext} rather than drained in a loop for two reasons. An + * {@code OnTimer} method may set further event time timers, and those must fire too: draining a + * snapshot silently truncates timer chains, the failure mode recorded for the RDD based runner + * in BEAM-12712. And firing one + * at a time keeps a timer heavy key from having to buffer all of its output at once. + * + *

    The clocks only need advancing once: {@code removeNext*} reads them live. As both clocks + * are pinned at {@code TIMESTAMP_MAX_VALUE}, a processing time timer re-armed from {@code + * OnTimer} targets a time past the pinned clock and never becomes eligible; batch mode makes no + * processing time guarantees (the RDD based runner drains the same way). + */ + private boolean fireNextTimer() throws Exception { + if (!clocksAdvanced) { + clocksAdvanced = true; + timerInternals.advanceInputWatermark(BoundedWindow.TIMESTAMP_MAX_VALUE); + timerInternals.advanceProcessingTime(BoundedWindow.TIMESTAMP_MAX_VALUE); + timerInternals.advanceSynchronizedProcessingTime(BoundedWindow.TIMESTAMP_MAX_VALUE); + } + TimerData timer = nextTimer(); + if (timer == null) { + return false; + } + fire(timer); + return true; + } + + private @Nullable TimerData nextTimer() { + TimerData timer = timerInternals.removeNextEventTimer(); + if (timer == null) { + timer = timerInternals.removeNextProcessingTimer(); + } + if (timer == null) { + timer = timerInternals.removeNextSynchronizedProcessingTimer(); + } + return timer; + } + + private void fire(TimerData timer) { + BoundedWindow window = + ((StateNamespaces.WindowNamespace) timer.getNamespace()).getWindow(); + runner.onTimer( + timer.getTimerId(), + timer.getTimerFamilyId(), + key, + window, + timer.getTimestamp(), + timer.getOutputTimestamp(), + timer.getDomain(), + CausedByDrain.NORMAL); + } + } +} diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/StatefulParDoTranslatorBatch.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/StatefulParDoTranslatorBatch.java new file mode 100644 index 000000000000..037595933d55 --- /dev/null +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/StatefulParDoTranslatorBatch.java @@ -0,0 +1,316 @@ +/* + * 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.spark.structuredstreaming.translation.batch; + +import static org.apache.beam.runners.spark.structuredstreaming.translation.helpers.EncoderHelpers.oneOfEncoder; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; +import static org.apache.spark.sql.functions.col; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.Map; +import org.apache.beam.runners.core.SideInputReader; +import org.apache.beam.runners.spark.SparkCommonPipelineOptions; +import org.apache.beam.runners.spark.structuredstreaming.metrics.MetricsAccumulator; +import org.apache.beam.runners.spark.structuredstreaming.translation.TransformTranslator; +import org.apache.beam.runners.spark.structuredstreaming.translation.batch.functions.SideInputValues; +import org.apache.beam.runners.spark.structuredstreaming.translation.batch.functions.SparkSideInputReader; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.reflect.DoFnSignature; +import org.apache.beam.sdk.transforms.reflect.DoFnSignatures; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionTuple; +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.Maps; +import org.apache.spark.api.java.function.FlatMapGroupsFunction; +import org.apache.spark.broadcast.Broadcast; +import org.apache.spark.sql.Column; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Encoder; +import org.apache.spark.sql.KeyValueGroupedDataset; +import org.apache.spark.sql.TypedColumn; +import org.apache.spark.storage.StorageLevel; +import scala.Tuple2; + +/** + * Translator for a stateful {@link ParDo.MultiOutput}, or one requiring time sorted input. + * + *

    Selected by {@link PipelineTranslatorBatch} in place of {@link ParDoTranslatorBatch} when the + * {@link DoFn} uses state, uses timers, or is annotated with {@link DoFn.RequiresTimeSortedInput}; + * see {@link #appliesTo}. + * + *

    Unlike {@link ParDoTranslatorBatch} this translator never produces an {@code + * UnresolvedTranslation}: a stateful {@link DoFn} must not be fused with neighbouring {@link ParDo + * ParDos}, because the fused runner cannot drive timers. Resolving the input dataset via {@code + * Context#getDataset} breaks any pending fusion chain. + * + *

    Additional (tagged) outputs are encoded as one column per tag, as in {@link + * ParDoTranslatorBatch}. + */ +@SuppressWarnings({"rawtypes", "unchecked"}) +class StatefulParDoTranslatorBatch + extends TransformTranslator< + PCollection>, PCollectionTuple, ParDo.MultiOutput, OutputT>> { + + StatefulParDoTranslatorBatch() { + // A stateful ParDo introduces a shuffle to co-locate and order each key, so it contributes to + // plan complexity much like GroupByKey rather than like a plain ParDo. + super(0.2f); + } + + /** + * Whether {@code transform} must be translated by this translator rather than {@link + * ParDoTranslatorBatch}. + * + *

    Note {@link DoFn.RequiresTimeSortedInput} is tested independently of state: the SDK only + * treats state and timers as making a {@link DoFn} stateful, so a {@code DoFn} carrying only that + * annotation reaches the runner with neither signature flag set. + */ + static boolean appliesTo(ParDo.MultiOutput transform) { + DoFnSignature signature = DoFnSignatures.signatureForDoFn(transform.getFn()); + return signature.usesState() + || signature.usesTimers() + || signature.processElement().requiresTimeSortedInput(); + } + + @Override + protected boolean canTranslate(ParDo.MultiOutput, OutputT> transform) { + DoFn, OutputT> doFn = transform.getFn(); + DoFnSignature signature = DoFnSignatures.signatureForDoFn(doFn); + + checkState( + appliesTo(transform), + "Not a stateful or time sorted DoFn, should have been translated by %s: %s", + ParDoTranslatorBatch.class.getSimpleName(), + doFn); + + checkState( + isSupported(), + "Stateful and time sorted ParDo require Spark 3.4+ " + + "(KeyValueGroupedDataset#flatMapSortedGroups): %s", + doFn); + + checkState( + !signature.processElement().isSplittable(), + "Not expected to directly translate splittable DoFn, should have been overridden: %s", + doFn); + + // Not implemented: firing @OnWindowExpiration requires tracking the windows observed per key + // and a dedicated firing pass at the end of each key, see + // https://github.com/apache/beam/issues/22524 + checkState( + signature.onWindowExpiration() == null, "onWindowExpiration is not supported: %s", doFn); + + SparkSideInputReader.validateMaterializations(transform.getSideInputs().values()); + return true; + } + + @Override + protected void translate(ParDo.MultiOutput, OutputT> transform, Context cxt) + throws IOException { + PCollection> input = (PCollection>) cxt.getInput(); + + validateKeyCoder(input.getCoder(), transform.getFn()); + validateWindowingStrategy(input.getWindowingStrategy(), transform.getFn()); + + TupleTag mainOut = transform.getMainOutputTag(); + // Filter out obsolete PCollections to only cache when absolutely necessary + Map, PCollection> outputs = + ParDoTranslatorBatch.skipUnconsumedOutputs( + cxt.getOutputs(), mainOut, transform.getAdditionalOutputTags(), cxt); + + KvCoder inputCoder = (KvCoder) input.getCoder(); + Encoder keyEnc = cxt.keyEncoderOf(inputCoder); + MetricsAccumulator metrics = MetricsAccumulator.getInstance(cxt.getSparkSession()); + SideInputReader sideInputReader = createSideInputReader(transform, cxt); + + // Group by key, then order each group by event time before handing it to the DoFn. The + // timestamp is a top level LongType column of the WindowedValue encoder (epoch millis), so + // ordering is plain signed numeric ordering; no composite sort key is needed. Nulls sort + // last: a null timestamp encodes END_OF_WINDOW (see GroupByKeyTranslatorBatch), which no + // concrete timestamp of the same window can exceed. Only null and concrete timestamps of + // different windows mixed into one key group may still order imprecisely; deriving the + // timestamp from the window column is not portable across Spark versions. + Column[] sortCols = new Column[] {col(TIMESTAMP_COLUMN).asc_nulls_last()}; + + if (outputs.size() > 1) { + // In case of multiple outputs / tags, map each tag to a column by index. + // At the end split the result into multiple datasets selecting one column each. + Map tagColIdx = ParDoTranslatorBatch.tagsColumnIndex(outputs.keySet()); + List>> encoders = createEncoders(outputs, tagColIdx, cxt); + + DoFnRunnerFactory, OutputT> runnerFactory = + DoFnRunnerFactory.simple(cxt.getCurrentTransform(), input, sideInputReader, false); + StatefulDoFnGroupFunction, Tuple2>> groupFn = + StatefulDoFnGroupFunction.multiOutput( + cxt.getOptionsSupplier(), metrics, runnerFactory, tagColIdx); + + SparkCommonPipelineOptions opts = cxt.getOptions().as(SparkCommonPipelineOptions.class); + StorageLevel storageLevel = StorageLevel.fromString(opts.getStorageLevel()); + + // Persist as wide rows with one column per TupleTag to support different schemas + Dataset>> allTagsDS = + cxt.getDataset(input) + .groupByKey(GroupByKeyHelpers.valueKey(), keyEnc) + .flatMapSortedGroups(sortCols, groupFn, oneOfEncoder(encoders)); + allTagsDS.persist(storageLevel); + + // divide into separate output datasets per tag + for (TupleTag tag : outputs.keySet()) { + int colIdx = checkStateNotNull(tagColIdx.get(tag.getId()), "Unknown tag"); + // Resolve specific column matching the tuple tag (by id) + TypedColumn>, WindowedValue> col = + (TypedColumn) col(Integer.toString(colIdx)).as(encoders.get(colIdx)); + + // Caching of the returned outputs is disabled to avoid caching the same data twice. + cxt.putDataset( + cxt.getOutput((TupleTag) tag), allTagsDS.filter(col.isNotNull()).select(col), false); + } + } else { + PCollection output = cxt.getOutput(mainOut); + // Obsolete outputs might have to be filtered out + boolean filterMainOutput = cxt.getOutputs().size() > 1; + DoFnRunnerFactory, OutputT> runnerFactory = + DoFnRunnerFactory.simple( + cxt.getCurrentTransform(), input, sideInputReader, filterMainOutput); + StatefulDoFnGroupFunction, WindowedValue> groupFn = + StatefulDoFnGroupFunction.singleOutput(cxt.getOptionsSupplier(), metrics, runnerFactory); + + Dataset> result = + cxt.getDataset(input) + .groupByKey(GroupByKeyHelpers.valueKey(), keyEnc) + .flatMapSortedGroups(sortCols, groupFn, cxt.windowedEncoder(output.getCoder())); + + cxt.putDataset(output, result); + } + } + + /** List of encoders matching the order of tagIds. */ + private List>> createEncoders( + Map, PCollection> outputs, Map tagIdColIdx, Context ctx) { + ArrayList>> encoders = new ArrayList<>(outputs.size()); + for (Map.Entry, PCollection> e : outputs.entrySet()) { + Encoder> enc = ctx.windowedEncoder((Coder) e.getValue().getCoder()); + int colIdx = checkStateNotNull(tagIdColIdx.get(e.getKey().getId())); + encoders.add(colIdx, enc); + } + return encoders; + } + + /** Field of the {@code WindowedValue} encoder holding the event time, as epoch millis. */ + private static final String TIMESTAMP_COLUMN = "timestamp"; + + private static final boolean SORTED_GROUPS_API_AVAILABLE = sortedGroupsApiAvailable(); + + /** + * Whether this Spark version supports stateful / time sorted ParDo: the required {@code + * flatMapSortedGroups(Column[], FlatMapGroupsFunction, Encoder)} only exists since Spark 3.4. + */ + static boolean isSupported() { + return SORTED_GROUPS_API_AVAILABLE; + } + + private static boolean sortedGroupsApiAvailable() { + try { + KeyValueGroupedDataset.class.getMethod( + "flatMapSortedGroups", Column[].class, FlatMapGroupsFunction.class, Encoder.class); + return true; + } catch (NoSuchMethodException e) { + return false; + } + } + + private SideInputReader createSideInputReader( + ParDo.MultiOutput, OutputT> transform, Context cxt) { + Collection> views = transform.getSideInputs().values(); + if (views.isEmpty()) { + return SparkSideInputReader.empty(); + } + Map>> broadcasts = + Maps.newHashMapWithExpectedSize(views.size()); + for (PCollectionView view : views) { + PCollection pCol = checkStateNotNull(view.getPCollection()); + Broadcast> broadcast = + (Broadcast) cxt.getSideInputBroadcast(pCol, SideInputValues.loader((PCollection) pCol)); + broadcasts.put(view.getTagInternal().getId(), broadcast); + } + return SparkSideInputReader.create(broadcasts); + } + + /** + * A stateful {@link DoFn} is keyed, and this translator co-locates and orders elements by the + * encoded key, so the key coder must be deterministic. + * + *

    {@code ParDo} already enforces both of these for {@code DoFns} using state or timers (see + * {@code ParDo.validateStateApplicableForInput}), but that validation is skipped for a {@link + * DoFn} carrying only {@link DoFn.RequiresTimeSortedInput}, which still reaches this translator. + * So it is checked here rather than assumed. + */ + @VisibleForTesting + static void validateKeyCoder(Coder coder, DoFn doFn) { + checkState( + coder instanceof KvCoder, + "Input to a stateful or time sorted ParDo requires a %s, but the coder was %s: %s", + KvCoder.class.getSimpleName(), + coder, + doFn); + + Coder keyCoder = ((KvCoder) coder).getKeyCoder(); + try { + keyCoder.verifyDeterministic(); + } catch (Coder.NonDeterministicException e) { + throw new IllegalStateException( + String.format( + "Input to a stateful or time sorted ParDo requires a deterministic key coder, " + + "but %s is not deterministic: %s", + keyCoder, doFn), + e); + } + } + + /** + * State is scoped per key and window, which is only well defined if windows are not still subject + * to merging. + * + *

    This deliberately mirrors Dataflow's {@code verifyStateSupportForWindowingStrategy} and + * tests {@link WindowingStrategy#needsMerge()} rather than {@code WindowFn#isNonMerging()}: after + * a {@link org.apache.beam.sdk.transforms.GroupByKey GroupByKey} the strategy keeps its merging + * {@code WindowFn} but is flagged as already merged, and such pipelines are legal. + */ + @VisibleForTesting + static void validateWindowingStrategy( + WindowingStrategy windowingStrategy, DoFn doFn) { + checkState( + !windowingStrategy.needsMerge(), + "Stateful and time sorted ParDo are not supported for merging windows, " + + "state cannot be scoped to a window that may still merge. WindowFn: %s, DoFn: %s", + windowingStrategy.getWindowFn(), + doFn); + } +} diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/StatefulParDoExecutionTest.java b/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/StatefulParDoExecutionTest.java new file mode 100644 index 000000000000..952f9c99a6b4 --- /dev/null +++ b/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/StatefulParDoExecutionTest.java @@ -0,0 +1,357 @@ +/* + * 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.spark.structuredstreaming.translation.batch; + +import java.io.Serializable; +import java.util.ArrayList; +import java.util.List; +import org.apache.beam.runners.spark.structuredstreaming.SparkSessionRule; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.coders.VarIntCoder; +import org.apache.beam.sdk.state.StateSpec; +import org.apache.beam.sdk.state.StateSpecs; +import org.apache.beam.sdk.state.TimeDomain; +import org.apache.beam.sdk.state.Timer; +import org.apache.beam.sdk.state.TimerSpec; +import org.apache.beam.sdk.state.TimerSpecs; +import org.apache.beam.sdk.state.ValueState; +import org.apache.beam.sdk.testing.PAssert; +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.ParDo; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionTuple; +import org.apache.beam.sdk.values.TimestampedValue; +import org.apache.beam.sdk.values.TupleTag; +import org.apache.beam.sdk.values.TupleTagList; +import org.joda.time.Duration; +import org.joda.time.Instant; +import org.junit.Assume; +import org.junit.BeforeClass; +import org.junit.ClassRule; +import org.junit.Rule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * Execution tests for {@link StatefulParDoTranslatorBatch} / {@link StatefulDoFnGroupFunction}: + * these run full pipelines, unlike {@link StatefulParDoTranslatorBatchTest} which only covers + * dispatch and translation preconditions. + */ +@RunWith(JUnit4.class) +public class StatefulParDoExecutionTest implements Serializable { + @ClassRule public static final SparkSessionRule SESSION = new SparkSessionRule(); + + @BeforeClass + public static void requireSortedGroupsApi() { + Assume.assumeTrue( + "Stateful ParDo requires Spark 3.4+", StatefulParDoTranslatorBatch.isSupported()); + } + + @Rule + public transient TestPipeline pipeline = + TestPipeline.fromOptions(SESSION.createPipelineOptions()); + + /** {@link ValueState} accumulates per key: totals must be scoped to the key, not shared. */ + @Test + public void testStatefulAccumulationPerKey() { + PCollection> result = + pipeline + .apply( + Create.timestamped( + TimestampedValue.of(KV.of("a", 1), sec(1)), + TimestampedValue.of(KV.of("a", 2), sec(2)), + TimestampedValue.of(KV.of("a", 3), sec(3)), + TimestampedValue.of(KV.of("b", 10), sec(1)), + TimestampedValue.of(KV.of("b", 20), sec(2)))) + .apply(ParDo.of(new RunningSumDoFn())); + + PAssert.that(result) + .containsInAnyOrder( + KV.of("a", 1), KV.of("a", 3), KV.of("a", 6), KV.of("b", 10), KV.of("b", 30)); + pipeline.run(); + } + + /** + * Many keys share the two shuffle partitions of the {@code local[2]} session, so several keys are + * served in sequence by the same DoFn instance and {@code MutableStepContext}. Each key's sums + * must be independent; any state bleeding between keys corrupts them. + */ + @Test + public void testStateIsolationAcrossManyKeysInOnePartition() { + List>> input = new ArrayList<>(); + List> expected = new ArrayList<>(); + for (int i = 0; i < 60; i++) { + String key = "key-" + i; + input.add(TimestampedValue.of(KV.of(key, i), sec(1))); + input.add(TimestampedValue.of(KV.of(key, 1000 + i), sec(2))); + expected.add(KV.of(key, i)); + expected.add(KV.of(key, 1000 + 2 * i)); + } + + PCollection> result = + pipeline.apply(Create.timestamped(input)).apply(ParDo.of(new RunningSumDoFn())); + + PAssert.that(result).containsInAnyOrder(expected); + pipeline.run(); + } + + /** + * Elements are created out of order but must reach a {@link DoFn.RequiresTimeSortedInput} DoFn in + * ascending timestamp order. The DoFn appends each value to state and emits the sequence so far, + * so the multiset of outputs pins the exact observation order. + */ + @Test + public void testRequiresTimeSortedInput() { + PCollection result = + pipeline + .apply( + Create.timestamped( + TimestampedValue.of(KV.of("k", 4), sec(4)), + TimestampedValue.of(KV.of("k", 1), sec(1)), + TimestampedValue.of(KV.of("k", 6), sec(6)), + TimestampedValue.of(KV.of("k", 3), sec(3)), + TimestampedValue.of(KV.of("k", 2), sec(2)), + TimestampedValue.of(KV.of("k", 5), sec(5)))) + .apply(ParDo.of(new TimeSortedSequenceDoFn())); + + PAssert.that(result) + .containsInAnyOrder("1", "1,2", "1,2,3", "1,2,3,4", "1,2,3,4,5", "1,2,3,4,5,6"); + pipeline.run(); + } + + /** An event time timer set in {@code @ProcessElement} must fire its {@code @OnTimer}. */ + @Test + public void testEventTimeTimerFires() { + PCollection result = + pipeline + .apply( + Create.timestamped( + TimestampedValue.of(KV.of("k", 1), sec(1)), + TimestampedValue.of(KV.of("k", 2), sec(2)))) + .apply(ParDo.of(new EventTimeTimerDoFn())); + + PAssert.that(result).containsInAnyOrder("elem-1", "elem-2", "timer-fired"); + pipeline.run(); + } + + /** + * An {@code @OnTimer} that re-sets its own timer must see every iteration fire: draining a + * snapshot of pending timers silently truncates such chains (the failure mode recorded for the + * RDD based runner in https://issues.apache.org/jira/browse/BEAM-12712). + */ + @Test + public void testLoopingTimerFiresAllIterations() { + PCollection result = + pipeline + .apply(Create.timestamped(TimestampedValue.of(KV.of("k", 1), sec(1)))) + .apply(ParDo.of(new LoopingTimerDoFn())); + + PAssert.that(result).containsInAnyOrder("fire-1", "fire-2", "fire-3", "fire-4", "fire-5"); + pipeline.run(); + } + + /** + * Timers fire while the bundle is still open: a buffer flushed by {@code @FinishBundle} must + * contain the {@code @OnTimer} contribution. Running {@code finishBundle} before the timers + * silently drops the timer's data. + */ + @Test + public void testFinishBundleFlushesTimerOutput() { + PCollection result = + pipeline + .apply( + Create.timestamped( + TimestampedValue.of(KV.of("k", 1), sec(1)), + TimestampedValue.of(KV.of("k", 2), sec(2)))) + .apply(ParDo.of(new BufferUntilFinishBundleDoFn())); + + PAssert.that(result).containsInAnyOrder("elem-1", "elem-2", "timer"); + pipeline.run(); + } + + /** + * A stateful {@link DoFn} with additional (tagged) outputs: per element the running sum goes to + * the main output while even values are also emitted to the additional output. + */ + @Test + public void testTaggedAdditionalOutput() { + TupleTag> sums = new TupleTag>() {}; + TupleTag evens = new TupleTag() {}; + + PCollectionTuple result = + pipeline + .apply( + Create.timestamped( + TimestampedValue.of(KV.of("a", 1), sec(1)), + TimestampedValue.of(KV.of("a", 2), sec(2)), + TimestampedValue.of(KV.of("b", 4), sec(1)))) + .apply( + ParDo.of(new RunningSumWithEvensDoFn(evens)) + .withOutputTags(sums, TupleTagList.of(evens))); + + PAssert.that(result.get(sums)).containsInAnyOrder(KV.of("a", 1), KV.of("a", 3), KV.of("b", 4)); + PAssert.that(result.get(evens)).containsInAnyOrder(2, 4); + pipeline.run(); + } + + private static Instant sec(long seconds) { + return new Instant(seconds * 1000); + } + + /** Emits the per key running sum for every element. */ + private static class RunningSumDoFn extends DoFn, KV> { + @StateId("sum") + private final StateSpec> sumSpec = StateSpecs.value(VarIntCoder.of()); + + @ProcessElement + public void processElement(ProcessContext c, @StateId("sum") ValueState sum) { + Integer current = sum.read(); + int newSum = (current == null ? 0 : current) + c.element().getValue(); + sum.write(newSum); + c.output(KV.of(c.element().getKey(), newSum)); + } + } + + /** Emits the per key running sum to the main output and even values to {@code evens}. */ + private static class RunningSumWithEvensDoFn + extends DoFn, KV> { + private final TupleTag evens; + + @StateId("sum") + private final StateSpec> sumSpec = StateSpecs.value(VarIntCoder.of()); + + RunningSumWithEvensDoFn(TupleTag evens) { + this.evens = evens; + } + + @ProcessElement + public void processElement(ProcessContext c, @StateId("sum") ValueState sum) { + Integer current = sum.read(); + int value = c.element().getValue(); + int newSum = (current == null ? 0 : current) + value; + sum.write(newSum); + c.output(KV.of(c.element().getKey(), newSum)); + if (value % 2 == 0) { + c.output(evens, value); + } + } + } + + /** Appends each value to state and emits the sequence observed so far. */ + private static class TimeSortedSequenceDoFn extends DoFn, String> { + @StateId("seen") + private final StateSpec> seenSpec = StateSpecs.value(StringUtf8Coder.of()); + + @RequiresTimeSortedInput + @ProcessElement + public void processElement(ProcessContext c, @StateId("seen") ValueState seen) { + String previous = seen.read(); + String sequence = + previous == null + ? c.element().getValue().toString() + : previous + "," + c.element().getValue(); + seen.write(sequence); + c.output(sequence); + } + } + + /** Sets one event time timer (re-set by each element, so it fires once). */ + private static class EventTimeTimerDoFn extends DoFn, String> { + @TimerId("timer") + private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.EVENT_TIME); + + @ProcessElement + public void processElement(ProcessContext c, @TimerId("timer") Timer timer) { + c.output("elem-" + c.element().getValue()); + timer.set(c.timestamp().plus(Duration.standardSeconds(10))); + } + + @OnTimer("timer") + public void onTimer(OnTimerContext c) { + c.output("timer-fired"); + } + } + + /** A bounded looping timer: each firing re-sets the timer until five have fired. */ + private static class LoopingTimerDoFn extends DoFn, String> { + @TimerId("loop") + private final TimerSpec loopSpec = TimerSpecs.timer(TimeDomain.EVENT_TIME); + + @StateId("fires") + private final StateSpec> firesSpec = StateSpecs.value(VarIntCoder.of()); + + @ProcessElement + public void processElement(ProcessContext c, @TimerId("loop") Timer loop) { + loop.set(c.timestamp().plus(Duration.standardSeconds(1))); + } + + @OnTimer("loop") + public void onTimer( + OnTimerContext c, + @TimerId("loop") Timer loop, + @StateId("fires") ValueState fires) { + Integer current = fires.read(); + int fired = (current == null ? 0 : current) + 1; + fires.write(fired); + c.output("fire-" + fired); + if (fired < 5) { + loop.set(c.fireTimestamp().plus(Duration.standardSeconds(1))); + } + } + } + + /** + * Buffers in the instance across {@code @ProcessElement} and {@code @OnTimer} and only outputs + * from {@code @FinishBundle}. + */ + private static class BufferUntilFinishBundleDoFn extends DoFn, String> { + @TimerId("flush") + private final TimerSpec flushSpec = TimerSpecs.timer(TimeDomain.EVENT_TIME); + + private transient List buffer; + + @StartBundle + public void startBundle() { + buffer = new ArrayList<>(); + } + + @ProcessElement + public void processElement(ProcessContext c, @TimerId("flush") Timer flush) { + buffer.add("elem-" + c.element().getValue()); + flush.set(c.timestamp().plus(Duration.standardSeconds(10))); + } + + @OnTimer("flush") + public void onTimer() { + buffer.add("timer"); + } + + @FinishBundle + public void finishBundle(FinishBundleContext c) { + for (String value : buffer) { + c.output(value, GlobalWindow.INSTANCE.maxTimestamp(), GlobalWindow.INSTANCE); + } + buffer = new ArrayList<>(); + } + } +} diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/StatefulParDoTranslatorBatchTest.java b/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/StatefulParDoTranslatorBatchTest.java new file mode 100644 index 000000000000..59f32b2a90f8 --- /dev/null +++ b/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/StatefulParDoTranslatorBatchTest.java @@ -0,0 +1,261 @@ +/* + * 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.spark.structuredstreaming.translation.batch; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import org.apache.beam.runners.spark.structuredstreaming.translation.TransformTranslator; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.CustomCoder; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.coders.VarIntCoder; +import org.apache.beam.sdk.state.StateSpec; +import org.apache.beam.sdk.state.StateSpecs; +import org.apache.beam.sdk.state.TimeDomain; +import org.apache.beam.sdk.state.Timer; +import org.apache.beam.sdk.state.TimerSpec; +import org.apache.beam.sdk.state.TimerSpecs; +import org.apache.beam.sdk.state.ValueState; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.windowing.FixedWindows; +import org.apache.beam.sdk.transforms.windowing.Sessions; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.TupleTag; +import org.apache.beam.sdk.values.TupleTagList; +import org.apache.beam.sdk.values.WindowingStrategy; +import org.joda.time.Duration; +import org.junit.Assume; +import org.junit.BeforeClass; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * Tests dispatch to {@link StatefulParDoTranslatorBatch} and its translation preconditions. + * + *

    These deliberately avoid {@code TestPipeline}: the behaviour under test is translator + * selection and validation, both of which are decided before any Spark session exists. + */ +@RunWith(JUnit4.class) +public class StatefulParDoTranslatorBatchTest { + + private static final KvCoder KV_CODER = + KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()); + + @BeforeClass + public static void requireSortedGroupsApi() { + Assume.assumeTrue( + "Stateful ParDo requires Spark 3.4+", StatefulParDoTranslatorBatch.isSupported()); + } + + // -------------------------------------------------------------------------------------------- + // Dispatch + // -------------------------------------------------------------------------------------------- + + @Test + public void appliesToStatefulDoFn() { + assertTrue(StatefulParDoTranslatorBatch.appliesTo(multiOutput(new StatefulDoFn()))); + } + + @Test + public void appliesToTimerDoFn() { + assertTrue(StatefulParDoTranslatorBatch.appliesTo(multiOutput(new TimerDoFn()))); + } + + /** + * A {@link DoFn} carrying only {@link DoFn.RequiresTimeSortedInput} is not considered stateful by + * the SDK, so dispatch must test the annotation separately from {@code usesState}/{@code + * usesTimers}. + */ + @Test + public void appliesToTimeSortedOnlyDoFn() { + assertTrue(StatefulParDoTranslatorBatch.appliesTo(multiOutput(new TimeSortedOnlyDoFn()))); + } + + @Test + public void doesNotApplyToPlainDoFn() { + assertFalse(StatefulParDoTranslatorBatch.appliesTo(multiOutput(new PlainDoFn()))); + } + + @Test + public void registryRoutesStatefulDoFnToStatefulTranslator() { + TransformTranslator translator = + new PipelineTranslatorBatch().getTransformTranslator(multiOutput(new StatefulDoFn())); + assertTrue( + "Expected StatefulParDoTranslatorBatch but got " + translator, + translator instanceof StatefulParDoTranslatorBatch); + } + + @Test + public void registryRoutesPlainDoFnToParDoTranslator() { + TransformTranslator translator = + new PipelineTranslatorBatch().getTransformTranslator(multiOutput(new PlainDoFn())); + assertTrue( + "Expected ParDoTranslatorBatch but got " + translator, + translator instanceof ParDoTranslatorBatch); + } + + // -------------------------------------------------------------------------------------------- + // Windowing precondition + // -------------------------------------------------------------------------------------------- + + @Test + public void rejectsMergingWindows() { + WindowingStrategy merging = + WindowingStrategy.of(Sessions.withGapDuration(Duration.standardMinutes(1))); + + IllegalStateException thrown = + assertThrows( + IllegalStateException.class, + () -> + StatefulParDoTranslatorBatch.validateWindowingStrategy( + merging, new StatefulDoFn())); + assertTrue(thrown.getMessage(), thrown.getMessage().contains("merging windows")); + } + + /** + * After a {@code GroupByKey} the strategy keeps its merging {@link Sessions} {@code WindowFn} but + * is flagged as already merged. Such pipelines are legal, so the precondition must test {@code + * needsMerge()} rather than {@code WindowFn#isNonMerging()}. + */ + @Test + public void acceptsWindowsAlreadyMerged() { + WindowingStrategy alreadyMerged = + WindowingStrategy.of(Sessions.withGapDuration(Duration.standardMinutes(1))) + .withAlreadyMerged(true); + + assertFalse("precondition of this test", alreadyMerged.getWindowFn().isNonMerging()); + StatefulParDoTranslatorBatch.validateWindowingStrategy(alreadyMerged, new StatefulDoFn()); + } + + @Test + public void acceptsNonMergingWindows() { + StatefulParDoTranslatorBatch.validateWindowingStrategy( + WindowingStrategy.of(FixedWindows.of(Duration.standardMinutes(1))), new StatefulDoFn()); + } + + // -------------------------------------------------------------------------------------------- + // Key coder precondition + // -------------------------------------------------------------------------------------------- + + @Test + public void acceptsDeterministicKeyCoder() { + StatefulParDoTranslatorBatch.validateKeyCoder(KV_CODER, new StatefulDoFn()); + } + + @Test + public void rejectsNonKvCoder() { + IllegalStateException thrown = + assertThrows( + IllegalStateException.class, + () -> + StatefulParDoTranslatorBatch.validateKeyCoder( + StringUtf8Coder.of(), new TimeSortedOnlyDoFn())); + assertTrue(thrown.getMessage(), thrown.getMessage().contains("KvCoder")); + } + + /** + * {@code ParDo} only validates the key coder for {@code DoFns} using state or timers, so a time + * sorted only {@code DoFn} can reach the translator with a non-deterministic key coder. + */ + @Test + public void rejectsNonDeterministicKeyCoder() { + Coder> nonDeterministic = + KvCoder.of(new NonDeterministicStringCoder(), VarIntCoder.of()); + + IllegalStateException thrown = + assertThrows( + IllegalStateException.class, + () -> + StatefulParDoTranslatorBatch.validateKeyCoder( + nonDeterministic, new TimeSortedOnlyDoFn())); + assertTrue(thrown.getMessage(), thrown.getMessage().contains("deterministic")); + } + + // -------------------------------------------------------------------------------------------- + // Fixtures + // -------------------------------------------------------------------------------------------- + + private static ParDo.MultiOutput, Integer> multiOutput( + DoFn, Integer> doFn) { + return ParDo.of(doFn).withOutputTags(new TupleTag() {}, TupleTagList.empty()); + } + + private static class PlainDoFn extends DoFn, Integer> { + @ProcessElement + public void processElement(ProcessContext ctx) { + ctx.output(ctx.element().getValue()); + } + } + + private static class StatefulDoFn extends DoFn, Integer> { + @StateId("value") + private final StateSpec> state = StateSpecs.value(VarIntCoder.of()); + + @ProcessElement + public void processElement(ProcessContext ctx, @StateId("value") ValueState state) { + ctx.output(ctx.element().getValue()); + } + } + + private static class TimerDoFn extends DoFn, Integer> { + @TimerId("timer") + private final TimerSpec timer = TimerSpecs.timer(TimeDomain.EVENT_TIME); + + @ProcessElement + public void processElement(ProcessContext ctx, @TimerId("timer") Timer timer) { + ctx.output(ctx.element().getValue()); + } + + @OnTimer("timer") + public void onTimer() {} + } + + private static class TimeSortedOnlyDoFn extends DoFn, Integer> { + @RequiresTimeSortedInput + @ProcessElement + public void processElement(ProcessContext ctx) { + ctx.output(ctx.element().getValue()); + } + } + + /** A String coder that refuses to declare itself deterministic. */ + private static class NonDeterministicStringCoder extends CustomCoder { + @Override + public void encode(String value, OutputStream outStream) throws IOException { + StringUtf8Coder.of().encode(value, outStream); + } + + @Override + public String decode(InputStream inStream) throws IOException { + return StringUtf8Coder.of().decode(inStream); + } + + @Override + public void verifyDeterministic() throws NonDeterministicException { + throw new NonDeterministicException(this, "not deterministic, by design, for this test"); + } + } +}