From cc6211bca1e21b0196f877d8a9e9e87de1c6d1f7 Mon Sep 17 00:00:00 2001 From: f-loris Date: Mon, 17 Aug 2026 20:55:16 +0200 Subject: [PATCH 1/2] [runners-spark] Support stateful ParDo in the Structured Streaming batch runner Stateful, timer-using and @RequiresTimeSortedInput DoFns are dispatched to a new StatefulParDoTranslatorBatch, which groups by key and sorts each group by event time; StatefulDoFnGroupFunction then runs the DoFn per key with in-memory state and timers. Additional (tagged) outputs are encoded as one column per tag and split into per-tag datasets, mirroring the stateless multi-output translation in ParDoTranslatorBatch, keeping the single-output fast path. Requires Spark 3.4+ for KeyValueGroupedDataset#flatMapSortedGroups; earlier versions are rejected at translation time rather than failing mid-job. - Enable the UsesStatefulParDo, UsesKeyInParDo, UsesTimersInParDo, UsesMapState, UsesMultimapState, UsesSetState and UsesTimerMap ValidatesRunner categories for validatesStructuredStreamingRunnerBatch - Exclude UsesOrderedListState: every such test also uses @OnWindowExpiration, which is not supported - Exclude PerKeyOrderingTest#testMultipleStatefulOrdering* by name: they build on PeriodicImpulse and so are unbounded, but are not categorized as UsesUnboundedPCollections, and adding that category upstream would also stop other runners running them @OnWindowExpiration remains unsupported (#22524). --- ...ostCommit_Java_ValidatesRunner_Spark4.json | 4 + ...idatesRunner_SparkStructuredStreaming.json | 3 +- CHANGES.md | 1 + runners/spark/spark_runner.gradle | 15 +- .../translation/batch/DoFnRunnerFactory.java | 26 +- .../batch/ParDoTranslatorBatch.java | 26 +- .../batch/PipelineTranslatorBatch.java | 16 + .../batch/StatefulDoFnGroupFunction.java | 392 ++++++++++++++++++ .../batch/StatefulParDoTranslatorBatch.java | 316 ++++++++++++++ .../batch/StatefulParDoExecutionTest.java | 357 ++++++++++++++++ .../StatefulParDoTranslatorBatchTest.java | 261 ++++++++++++ 11 files changed, 1396 insertions(+), 21 deletions(-) create mode 100644 .github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark4.json create mode 100644 runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/StatefulDoFnGroupFunction.java create mode 100644 runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/StatefulParDoTranslatorBatch.java create mode 100644 runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/StatefulParDoExecutionTest.java create mode 100644 runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/StatefulParDoTranslatorBatchTest.java 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"); + } + } +} From 0bd5d68d2ba03d49d641c9fd876fe0dad770f285 Mon Sep 17 00:00:00 2001 From: f-loris Date: Thu, 20 Aug 2026 11:38:04 +0200 Subject: [PATCH 2/2] [runners-spark] Address review comments on stateful ParDo translation - Drop the class level @SuppressWarnings from StatefulParDoTranslatorBatch. Its raw typed createSideInputReader is gone, replaced by the shared static one in ParDoTranslatorBatch, which now takes a TranslationState. - Index createEncoders by column rather than appending at an index into an empty ArrayList, which only worked while the iteration order of the outputs happened to match the column order. Shared as a single static between both translators. - Note at the dispatch site why ParDo.MultiOutput is resolved by predicate ahead of the class keyed registry. - Drop the redundant @NonNull bound in StatefulDoFnGroupFunction. --- .../batch/ParDoTranslatorBatch.java | 24 ++++++----- .../batch/PipelineTranslatorBatch.java | 3 ++ .../batch/StatefulDoFnGroupFunction.java | 3 +- .../batch/StatefulParDoTranslatorBatch.java | 42 ++----------------- 4 files changed, 22 insertions(+), 50 deletions(-) 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 cb05da63dc3e..14058a73733d 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 @@ -24,7 +24,7 @@ import static org.apache.spark.sql.functions.col; import java.io.IOException; -import java.util.ArrayList; +import java.util.Arrays; import java.util.Collection; import java.util.Collections; import java.util.List; @@ -247,20 +247,24 @@ static Map tagsColumnIndex(Collection> tags) { return index; } - /** List of encoders matching the order of tagIds. */ - private List>> createEncoders( - Map, PCollection> outputs, Map tagIdColIdx, Context ctx) { - ArrayList>> encoders = new ArrayList<>(outputs.size()); + /** List of encoders indexed by column index, as assigned by {@code tagIdColIdx}. */ + @SuppressWarnings("rawtypes") // generic array creation + static List>> createEncoders( + Map, PCollection> outputs, + Map tagIdColIdx, + TransformTranslator.Context ctx) { + // Indexed rather than appended, so the iteration order of outputs need not match the columns. + Encoder>[] encoders = new Encoder[outputs.size()]; for (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); + encoders[colIdx] = ctx.windowedEncoder((Coder) e.getValue().getCoder()); } - return encoders; + return Arrays.asList(encoders); } - private SideInputReader createSideInputReader( - Collection> views, Context cxt) { + /** Broadcasts {@code views}, if any, and exposes them as a {@link SideInputReader}. */ + static SideInputReader createSideInputReader( + Collection> views, TranslationState cxt) { if (views.isEmpty()) { return SparkSideInputReader.empty(); } 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 d6d85d977c04..ba7cbb0fa035 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 @@ -99,6 +99,9 @@ public class PipelineTranslatorBatch extends PipelineTranslator { @Nullable protected > TransformTranslator getTransformTranslator(TransformT transform) { + // Resolved ahead of the class keyed registry: ParDo.MultiOutput maps to a different translator + // depending on the DoFn signature, which a lookup by transform class alone cannot express. This + // is the predicated dispatch of the TODO above, limited to the single transform needing it. if (transform instanceof ParDo.MultiOutput && StatefulParDoTranslatorBatch.appliesTo((ParDo.MultiOutput) transform)) { return STATEFUL_PARDO_TRANSLATOR; 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 index 77eb0a7cb65d..b5c2d602407e 100644 --- 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 @@ -46,7 +46,6 @@ 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; @@ -71,7 +70,7 @@ * 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> +abstract class StatefulDoFnGroupFunction, OutT> implements FlatMapGroupsFunction, OutT> { private final Supplier options; 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 index 037595933d55..75d84b630bb8 100644 --- 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 @@ -23,15 +23,12 @@ 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; @@ -42,14 +39,11 @@ 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; @@ -73,7 +67,6 @@ *

    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>> { @@ -148,7 +141,8 @@ protected void translate(ParDo.MultiOutput, OutputT> transform, Context KvCoder inputCoder = (KvCoder) input.getCoder(); Encoder keyEnc = cxt.keyEncoderOf(inputCoder); MetricsAccumulator metrics = MetricsAccumulator.getInstance(cxt.getSparkSession()); - SideInputReader sideInputReader = createSideInputReader(transform, cxt); + SideInputReader sideInputReader = + ParDoTranslatorBatch.createSideInputReader(transform.getSideInputs().values(), 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 @@ -163,7 +157,8 @@ protected void translate(ParDo.MultiOutput, OutputT> transform, Context // 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); + List>> encoders = + ParDoTranslatorBatch.createEncoders(outputs, tagColIdx, cxt); DoFnRunnerFactory, OutputT> runnerFactory = DoFnRunnerFactory.simple(cxt.getCurrentTransform(), input, sideInputReader, false); @@ -211,18 +206,6 @@ protected void translate(ParDo.MultiOutput, OutputT> transform, Context } } - /** 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"; @@ -246,23 +229,6 @@ private static boolean sortedGroupsApiAvailable() { } } - 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.