This is an automated email from the ASF dual-hosted git repository.
Abacn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 67018be2650 [runners-spark] Support stateful ParDo in the Structured
Streaming batch runner (#39793)
67018be2650 is described below
commit 67018be26507a4afdda3a8368c53db4a0ddf0366
Author: Fabian Loris <[email protected]>
AuthorDate: Thu Aug 20 22:41:44 2026 +0200
[runners-spark] Support stateful ParDo in the Structured Streaming batch
runner (#39793)
* [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).
---
...eam_PostCommit_Java_ValidatesRunner_Spark4.json | 4 +
...a_ValidatesRunner_SparkStructuredStreaming.json | 3 +-
CHANGES.md | 1 +
runners/spark/spark_runner.gradle | 15 +-
.../translation/batch/DoFnRunnerFactory.java | 26 +-
.../translation/batch/ParDoTranslatorBatch.java | 50 +--
.../translation/batch/PipelineTranslatorBatch.java | 19 +
.../batch/StatefulDoFnGroupFunction.java | 391 +++++++++++++++++++++
.../batch/StatefulParDoTranslatorBatch.java | 282 +++++++++++++++
.../batch/StatefulParDoExecutionTest.java | 357 +++++++++++++++++++
.../batch/StatefulParDoTranslatorBatchTest.java | 261 ++++++++++++++
11 files changed, 1378 insertions(+), 31 deletions(-)
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 00000000000..e3d6056a5de
--- /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 77f63217b86..cad8d98b8ea 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 73966a48313..56c062790e2 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 77da3d36db9..2a161db0a82 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 5e8703a05b0..ce4155ee8e1 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 java.util.Map;
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 @@ abstract class DoFnRunnerFactory<InT, T> implements
Serializable {
abstract DoFnRunnerWithTeardown<InT, T> create(
PipelineOptions options, MetricsAccumulator metrics,
WindowedValueMultiReceiver output);
+ /**
+ * Creates a runner backed by {@code stepContext} so that state and timers
are available.
+ *
+ * <p>Only supported for a single, unfused {@link DoFn}: a fused runner
cannot drive timers.
+ */
+ DoFnRunnerWithTeardown<InT, T> 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 @@ abstract class DoFnRunnerFactory<InT, T> implements
Serializable {
@Override
DoFnRunnerWithTeardown<InT, T> create(
PipelineOptions options, MetricsAccumulator metrics,
WindowedValueMultiReceiver output) {
+ return create(options, metrics, output, new NoOpStepContext());
+ }
+
+ @Override
+ DoFnRunnerWithTeardown<InT, T> create(
+ PipelineOptions options,
+ MetricsAccumulator metrics,
+ WindowedValueMultiReceiver output,
+ StepContext stepContext) {
DoFnRunner<InT, T> simpleRunner =
DoFnRunners.simpleRunner(
options,
@@ -136,7 +160,7 @@ abstract class DoFnRunnerFactory<InT, T> implements
Serializable {
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 0f43f329b0d..14058a73733 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.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Pr
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;
@@ -35,6 +35,7 @@ import org.apache.beam.runners.core.DoFnRunners;
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 @@ import scala.Tuple2;
*
* <p>Each tag is encoded as individual column with a respective schema &
encoder each.
*
+ * <p>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.
+ *
* <p>TODO:
- * <li>Add support for state and timers.
* <li>Add support for SplittableDoFn
*/
class ParDoTranslatorBatch<InputT, OutputT>
@@ -87,18 +91,18 @@ class ParDoTranslatorBatch<InputT, OutputT>
"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 @@ class ParDoTranslatorBatch<InputT, OutputT>
* <p>This can help to avoid unnecessary caching in case of multiple outputs
if only {@code
* mainTag} is consumed.
*/
- private Map<TupleTag<?>, PCollection<?>> skipUnconsumedOutputs(
+ static Map<TupleTag<?>, PCollection<?>> skipUnconsumedOutputs(
Map<TupleTag<?>, 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 @@ class ParDoTranslatorBatch<InputT, OutputT>
}
}
- private Map<String, Integer> tagsColumnIndex(Collection<TupleTag<?>> tags) {
+ static Map<String, Integer> tagsColumnIndex(Collection<TupleTag<?>> tags) {
Map<String, Integer> index = Maps.newHashMapWithExpectedSize(tags.size());
for (TupleTag<?> tag : tags) {
index.put(tag.getId(), index.size());
@@ -243,20 +247,24 @@ class ParDoTranslatorBatch<InputT, OutputT>
return index;
}
- /** List of encoders matching the order of tagIds. */
- private List<Encoder<WindowedValue<Object>>> createEncoders(
- Map<TupleTag<?>, PCollection<?>> outputs, Map<String, Integer>
tagIdColIdx, Context ctx) {
- ArrayList<Encoder<WindowedValue<Object>>> encoders = new
ArrayList<>(outputs.size());
+ /** List of encoders indexed by column index, as assigned by {@code
tagIdColIdx}. */
+ @SuppressWarnings("rawtypes") // generic array creation
+ static List<Encoder<WindowedValue<Object>>> createEncoders(
+ Map<TupleTag<?>, PCollection<?>> outputs,
+ Map<String, Integer> tagIdColIdx,
+ TransformTranslator<?, ?, ?>.Context ctx) {
+ // Indexed rather than appended, so the iteration order of outputs need
not match the columns.
+ Encoder<WindowedValue<Object>>[] encoders = new Encoder[outputs.size()];
for (Entry<TupleTag<?>, PCollection<?>> e : outputs.entrySet()) {
- Encoder<WindowedValue<Object>> 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 <T> SideInputReader createSideInputReader(
- Collection<PCollectionView<?>> views, Context cxt) {
+ /** Broadcasts {@code views}, if any, and exposes them as a {@link
SideInputReader}. */
+ static <T> SideInputReader createSideInputReader(
+ Collection<PCollectionView<?>> 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 c4a18801ccb..ba7cbb0fa03 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,30 @@ 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.
+ *
+ * <p>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 <InT extends PInput, OutT extends POutput, TransformT extends
PTransform<InT, OutT>>
TransformTranslator<InT, OutT, TransformT>
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;
+ }
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 00000000000..b5c2d602407
--- /dev/null
+++
b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/StatefulDoFnGroupFunction.java
@@ -0,0 +1,391 @@
+/*
+ * 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.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.
+ *
+ * <p>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.
+ *
+ * <p>The {@link DoFn} is set up <b>once per task</b> 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 <em>bundle</em>, which is the level the model does allow to vary, and
its own state and
+ * timers via {@link MutableStepContext}.
+ *
+ * <p>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<K, InT extends KV<K, ?>, OutT>
+ implements FlatMapGroupsFunction<K, WindowedValue<InT>, OutT> {
+
+ private final Supplier<PipelineOptions> options;
+ private final MetricsAccumulator metrics;
+ private final DoFnRunnerFactory<InT, ?> factory;
+
+ private transient @Nullable Deque<OutT> buffer;
+ private transient @Nullable MutableStepContext stepContext;
+ private transient @Nullable DoFnRunnerWithTeardown<InT, ?> doFnRunner;
+ private transient boolean needsBundleStart;
+ private transient boolean isTornDown;
+
+ private StatefulDoFnGroupFunction(
+ Supplier<PipelineOptions> options,
+ MetricsAccumulator metrics,
+ DoFnRunnerFactory<InT, ?> factory) {
+ this.options = options;
+ this.metrics = metrics;
+ this.factory = factory;
+ }
+
+ /**
+ * {@link StatefulDoFnGroupFunction} emitting a single output of type {@link
WindowedValue} of
+ * {@link FnOutT}.
+ */
+ static <K, InT extends KV<K, ?>, FnOutT>
+ StatefulDoFnGroupFunction<K, InT, WindowedValue<FnOutT>> singleOutput(
+ Supplier<PipelineOptions> options,
+ MetricsAccumulator metrics,
+ DoFnRunnerFactory<InT, FnOutT> 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 <K, InT extends KV<K, ?>, FnOutT, OutT>
+ StatefulDoFnGroupFunction<K, InT, Tuple2<Integer, WindowedValue<OutT>>>
multiOutput(
+ Supplier<PipelineOptions> options,
+ MetricsAccumulator metrics,
+ DoFnRunnerFactory<InT, FnOutT> factory,
+ Map<String, Integer> tagColIdx) {
+ return new MultiOut<>(options, metrics, factory, tagColIdx);
+ }
+
+ @Override
+ public Iterator<OutT> call(K key, Iterator<WindowedValue<InT>> values) {
+ DoFnRunnerWithTeardown<InT, ?> 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<InT, ?> runner() {
+ DoFnRunnerWithTeardown<InT, ?> runner = doFnRunner;
+ if (runner == null) {
+ MutableStepContext ctx = new MutableStepContext();
+ Deque<OutT> 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<OutT> buffer() {
+ Deque<OutT> buf = buffer;
+ if (buf == null) {
+ throw new IllegalStateException("Buffer requested before the runner was
created");
+ }
+ return buf;
+ }
+
+ private void teardownOnce() {
+ DoFnRunnerWithTeardown<InT, ?> 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<OutT> buffer);
+
+ /**
+ * {@link StatefulDoFnGroupFunction} emitting a single output of type {@link
WindowedValue} of
+ * {@link FnOutT}.
+ */
+ private static class SingleOut<K, InT extends KV<K, ?>, FnOutT>
+ extends StatefulDoFnGroupFunction<K, InT, WindowedValue<FnOutT>> {
+ private SingleOut(
+ Supplier<PipelineOptions> options,
+ MetricsAccumulator metrics,
+ DoFnRunnerFactory<InT, FnOutT> factory) {
+ super(options, metrics, factory);
+ }
+
+ @Override
+ WindowedValueMultiReceiver outputManager(Deque<WindowedValue<FnOutT>>
buffer) {
+ return new WindowedValueMultiReceiver() {
+ @Override
+ public <T> void output(TupleTag<T> tag, WindowedValue<T> output) {
+ buffer.add((WindowedValue<FnOutT>) 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<K, InT extends KV<K, ?>, FnOutT, OutT>
+ extends StatefulDoFnGroupFunction<K, InT, Tuple2<Integer,
WindowedValue<OutT>>> {
+ private final Map<String, Integer> tagColIdx;
+
+ private MultiOut(
+ Supplier<PipelineOptions> options,
+ MetricsAccumulator metrics,
+ DoFnRunnerFactory<InT, FnOutT> factory,
+ Map<String, Integer> tagColIdx) {
+ super(options, metrics, factory);
+ this.tagColIdx = tagColIdx;
+ }
+
+ @Override
+ WindowedValueMultiReceiver outputManager(Deque<Tuple2<Integer,
WindowedValue<OutT>>> buffer) {
+ return new WindowedValueMultiReceiver() {
+ @Override
+ public <T> void output(TupleTag<T> tag, WindowedValue<T> 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<OutT>) 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.
+ *
+ * <p>{@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<OutT> {
+ private final Iterator<WindowedValue<InT>> groupIt;
+ private final K key;
+ private final DoFnRunnerWithTeardown<InT, ?> runner;
+ private final InMemoryTimerInternals timerInternals;
+
+ private boolean areTimersDrained;
+ private boolean clocksAdvanced;
+ private boolean isBundleFinished;
+
+ private StatefulGroupIt(
+ K key, Iterator<WindowedValue<InT>> groupIt,
DoFnRunnerWithTeardown<InT, ?> runner) {
+ this.key = key;
+ this.groupIt = groupIt;
+ this.runner = runner;
+ this.timerInternals = stepContext().timers();
+ }
+
+ @Override
+ protected @CheckForNull OutT computeNext() {
+ Deque<OutT> 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.
+ *
+ * <p>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 <a
href="https://issues.apache.org/jira/browse/BEAM-12712";>BEAM-12712</a>. And
firing one
+ * at a time keeps a timer heavy key from having to buffer all of its
output at once.
+ *
+ * <p>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 00000000000..75d84b630bb
--- /dev/null
+++
b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/StatefulParDoTranslatorBatch.java
@@ -0,0 +1,282 @@
+/*
+ * 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.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.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.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.spark.api.java.function.FlatMapGroupsFunction;
+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.
+ *
+ * <p>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}.
+ *
+ * <p>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.
+ *
+ * <p>Additional (tagged) outputs are encoded as one column per tag, as in
{@link
+ * ParDoTranslatorBatch}.
+ */
+class StatefulParDoTranslatorBatch<K, V, OutputT>
+ extends TransformTranslator<
+ PCollection<? extends KV<K, V>>, PCollectionTuple,
ParDo.MultiOutput<KV<K, V>, 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}.
+ *
+ * <p>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<KV<K, V>, OutputT>
transform) {
+ DoFn<KV<K, V>, 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<KV<K, V>, OutputT> transform,
Context cxt)
+ throws IOException {
+ PCollection<KV<K, V>> input = (PCollection<KV<K, V>>) cxt.getInput();
+
+ validateKeyCoder(input.getCoder(), transform.getFn());
+ validateWindowingStrategy(input.getWindowingStrategy(), transform.getFn());
+
+ TupleTag<OutputT> mainOut = transform.getMainOutputTag();
+ // Filter out obsolete PCollections to only cache when absolutely necessary
+ Map<TupleTag<?>, PCollection<?>> outputs =
+ ParDoTranslatorBatch.skipUnconsumedOutputs(
+ cxt.getOutputs(), mainOut, transform.getAdditionalOutputTags(),
cxt);
+
+ KvCoder<K, V> inputCoder = (KvCoder<K, V>) input.getCoder();
+ Encoder<K> keyEnc = cxt.keyEncoderOf(inputCoder);
+ MetricsAccumulator metrics =
MetricsAccumulator.getInstance(cxt.getSparkSession());
+ 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
+ // 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<String, Integer> tagColIdx =
ParDoTranslatorBatch.tagsColumnIndex(outputs.keySet());
+ List<Encoder<WindowedValue<Object>>> encoders =
+ ParDoTranslatorBatch.createEncoders(outputs, tagColIdx, cxt);
+
+ DoFnRunnerFactory<KV<K, V>, OutputT> runnerFactory =
+ DoFnRunnerFactory.simple(cxt.getCurrentTransform(), input,
sideInputReader, false);
+ StatefulDoFnGroupFunction<K, KV<K, V>, Tuple2<Integer,
WindowedValue<Object>>> 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<Tuple2<Integer, WindowedValue<Object>>> 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<Tuple2<Integer, WindowedValue<Object>>,
WindowedValue<Object>> 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<OutputT> output = cxt.getOutput(mainOut);
+ // Obsolete outputs might have to be filtered out
+ boolean filterMainOutput = cxt.getOutputs().size() > 1;
+ DoFnRunnerFactory<KV<K, V>, OutputT> runnerFactory =
+ DoFnRunnerFactory.simple(
+ cxt.getCurrentTransform(), input, sideInputReader,
filterMainOutput);
+ StatefulDoFnGroupFunction<K, KV<K, V>, WindowedValue<OutputT>> groupFn =
+ StatefulDoFnGroupFunction.singleOutput(cxt.getOptionsSupplier(),
metrics, runnerFactory);
+
+ Dataset<WindowedValue<OutputT>> result =
+ cxt.getDataset(input)
+ .groupByKey(GroupByKeyHelpers.valueKey(), keyEnc)
+ .flatMapSortedGroups(sortCols, groupFn,
cxt.windowedEncoder(output.getCoder()));
+
+ cxt.putDataset(output, result);
+ }
+ }
+
+ /** 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;
+ }
+ }
+
+ /**
+ * 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.
+ *
+ * <p>{@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.
+ *
+ * <p>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 00000000000..952f9c99a6b
--- /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<KV<String, Integer>> 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<TimestampedValue<KV<String, Integer>>> input = new ArrayList<>();
+ List<KV<String, Integer>> 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<KV<String, Integer>> 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<String> 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<String> 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<String> 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<String> 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<KV<String, Integer>> sums = new TupleTag<KV<String, Integer>>()
{};
+ TupleTag<Integer> evens = new TupleTag<Integer>() {};
+
+ 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<String, Integer>,
KV<String, Integer>> {
+ @StateId("sum")
+ private final StateSpec<ValueState<Integer>> sumSpec =
StateSpecs.value(VarIntCoder.of());
+
+ @ProcessElement
+ public void processElement(ProcessContext c, @StateId("sum")
ValueState<Integer> 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<String, Integer>, KV<String, Integer>> {
+ private final TupleTag<Integer> evens;
+
+ @StateId("sum")
+ private final StateSpec<ValueState<Integer>> sumSpec =
StateSpecs.value(VarIntCoder.of());
+
+ RunningSumWithEvensDoFn(TupleTag<Integer> evens) {
+ this.evens = evens;
+ }
+
+ @ProcessElement
+ public void processElement(ProcessContext c, @StateId("sum")
ValueState<Integer> 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<KV<String,
Integer>, String> {
+ @StateId("seen")
+ private final StateSpec<ValueState<String>> seenSpec =
StateSpecs.value(StringUtf8Coder.of());
+
+ @RequiresTimeSortedInput
+ @ProcessElement
+ public void processElement(ProcessContext c, @StateId("seen")
ValueState<String> 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<KV<String, Integer>,
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<KV<String, Integer>,
String> {
+ @TimerId("loop")
+ private final TimerSpec loopSpec = TimerSpecs.timer(TimeDomain.EVENT_TIME);
+
+ @StateId("fires")
+ private final StateSpec<ValueState<Integer>> 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<Integer> 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<KV<String,
Integer>, String> {
+ @TimerId("flush")
+ private final TimerSpec flushSpec =
TimerSpecs.timer(TimeDomain.EVENT_TIME);
+
+ private transient List<String> 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 00000000000..59f32b2a90f
--- /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.
+ *
+ * <p>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<String, Integer> 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<KV<String, Integer>> 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<KV<String, Integer>, Integer> multiOutput(
+ DoFn<KV<String, Integer>, Integer> doFn) {
+ return ParDo.of(doFn).withOutputTags(new TupleTag<Integer>() {},
TupleTagList.empty());
+ }
+
+ private static class PlainDoFn extends DoFn<KV<String, Integer>, Integer> {
+ @ProcessElement
+ public void processElement(ProcessContext ctx) {
+ ctx.output(ctx.element().getValue());
+ }
+ }
+
+ private static class StatefulDoFn extends DoFn<KV<String, Integer>, Integer>
{
+ @StateId("value")
+ private final StateSpec<ValueState<Integer>> state =
StateSpecs.value(VarIntCoder.of());
+
+ @ProcessElement
+ public void processElement(ProcessContext ctx, @StateId("value")
ValueState<Integer> state) {
+ ctx.output(ctx.element().getValue());
+ }
+ }
+
+ private static class TimerDoFn extends DoFn<KV<String, Integer>, 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<KV<String, Integer>,
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<String>
{
+ @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");
+ }
+ }
+}
