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 e261c9f7689 [#36841] Spark 4 streaming: windowed GroupByKey and
Combine via transformWithState (#40703)
e261c9f7689 is described below
commit e261c9f7689dbf6f0c6ecca75a8c11dd0d5eb070
Author: Tobias Kaymak <[email protected]>
AuthorDate: Thu Oct 22 17:20:47 2026 +0200
[#36841] Spark 4 streaming: windowed GroupByKey and Combine via
transformWithState (#40703)
* [#36841] Spark 4 streaming: windowed GroupByKey and Combine via
transformWithState
* Address review: move taskRunner to the ParDo processor, one
SparkStateInternals per visit
---
...a_ValidatesRunner_SparkStructuredStreaming.json | 2 +-
.../translation/PipelineTranslatorStreaming.java | 51 +++-
...tor.java => GroupByKeyStreamingTranslator.java} | 42 +---
.../StatefulParDoStreamingTranslator.java | 6 +-
.../streaming/state/BeamStatefulProcessor.java | 173 ++++---------
.../state/GroupByKeyStatefulProcessor.java | 141 +++++++++++
.../streaming/state/ParDoStatefulProcessor.java | 156 ++++++++++++
.../PipelineTranslatorStreamingTest.java | 55 ++++-
.../streaming/StatefulParDoStreamingTest.java | 15 +-
.../translation/streaming/StreamingTestUtils.java | 13 +
.../streaming/WindowedGroupByKeyStreamingTest.java | 274 +++++++++++++++++++++
11 files changed, 733 insertions(+), 195 deletions(-)
diff --git
a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json
b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json
index c0d5a77072b..1b832d7f19c 100644
---
a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json
+++
b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json
@@ -12,5 +12,5 @@
"https://github.com/apache/beam/pull/40103": "noting that PR #40103 should
run this test",
"https://github.com/apache/beam/issues/40427": "Spark ValidatesRunner tests
run serially",
"https://github.com/apache/beam/pull/40281": "noting that PR #40281 should
run this test",
- "https://github.com/apache/beam/issues/36841": "Spark 4 streaming null keys"
+ "https://github.com/apache/beam/issues/36841": "Spark 4 streaming GroupByKey"
}
diff --git
a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java
b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java
index 5a18b5da6fd..6f2d14a69b8 100644
---
a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java
+++
b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java
@@ -20,6 +20,7 @@ package
org.apache.beam.runners.spark.structuredstreaming.translation;
import java.util.Collection;
import org.apache.beam.runners.spark.SparkCommonPipelineOptions;
import
org.apache.beam.runners.spark.structuredstreaming.translation.batch.PipelineTranslatorCommon;
+import
org.apache.beam.runners.spark.structuredstreaming.translation.streaming.GroupByKeyStreamingTranslator;
import
org.apache.beam.runners.spark.structuredstreaming.translation.streaming.ReadUnboundedTranslator;
import
org.apache.beam.runners.spark.structuredstreaming.translation.streaming.StatefulParDoStreamingTranslator;
import org.apache.beam.sdk.annotations.Internal;
@@ -33,20 +34,21 @@ import org.apache.beam.sdk.transforms.PTransform;
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.transforms.windowing.DefaultTrigger;
import org.apache.beam.sdk.util.construction.SplittableParDo;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PInput;
import org.apache.beam.sdk.values.POutput;
import org.apache.beam.sdk.values.WindowingStrategy;
+import org.apache.beam.sdk.values.WindowingStrategy.AccumulationMode;
import org.apache.spark.sql.SparkSession;
import org.checkerframework.checker.nullness.qual.Nullable;
/**
* Pipeline translator for streaming pipelines on Spark 4. It extends the
common registry to reuse
- * the stateless single output ParDo, Window.Assign, Flatten and Reshuffle
translators, which are
- * safe on a streaming Dataset. Every other primitive fails at translation,
the batch translators
- * for them persist or collect the Dataset, which Spark rejects for streaming
plans.
+ * the stateless single output ParDo, Window.Assign, Flatten, Reshuffle and
Combine.GroupedValues
+ * translators, which are safe on a streaming Dataset. Unsupported primitives
fail at translation.
*/
@Internal
public class PipelineTranslatorStreaming extends PipelineTranslatorCommon {
@@ -61,7 +63,18 @@ public class PipelineTranslatorStreaming extends
PipelineTranslatorCommon {
protected void translate(
ParDo.MultiOutput<KV<Object, Object>, Object> transform, Context
cxt) {
PCollection<?> input = (PCollection<?>) cxt.getInput();
- checkNotMergingWindows(input.getWindowingStrategy());
+ checkNotMergingWindows(input.getWindowingStrategy(), "Stateful
ParDo");
+ super.translate(transform, cxt);
+ }
+ };
+
+ @SuppressWarnings({"rawtypes", "unchecked"})
+ private static final TransformTranslator GROUP_BY_KEY =
+ new GroupByKeyStreamingTranslator<Object, Object>() {
+ @Override
+ protected void translate(GroupByKey<Object, Object> transform, Context
cxt) {
+ PCollection<?> input = (PCollection<?>) cxt.getInput();
+ checkSupportedGroupByKey(input);
super.translate(transform, cxt);
}
};
@@ -92,11 +105,17 @@ public class PipelineTranslatorStreaming extends
PipelineTranslatorCommon {
}
if (transform instanceof GroupByKey) {
- throw new UnsupportedOperationException("GroupByKey" + NOT_SUPPORTED);
+ return GROUP_BY_KEY;
+ }
+
+ if (transform instanceof Combine.Globally
+ && ((Combine.Globally<?, ?>) transform).isInsertDefault()) {
+ throw new UnsupportedOperationException(
+ "Combine.globally() with a default value (use withoutDefaults())" +
NOT_SUPPORTED);
}
- if (transform instanceof Combine.PerKey) {
- throw new UnsupportedOperationException("Combine.perKey" +
NOT_SUPPORTED);
+ if (transform instanceof Combine.PerKey || transform instanceof
Combine.Globally) {
+ return null;
}
if (transform instanceof ParDo.MultiOutput) {
@@ -153,10 +172,22 @@ public class PipelineTranslatorStreaming extends
PipelineTranslatorCommon {
}
}
- private static void checkNotMergingWindows(WindowingStrategy<?, ?>
windowingStrategy) {
- if (windowingStrategy.needsMerge()) {
+ private static void checkSupportedGroupByKey(PCollection<?> input) {
+ WindowingStrategy<?, ?> strategy = input.getWindowingStrategy();
+ checkNotMergingWindows(strategy, "GroupByKey");
+ if (!(strategy.getTrigger() instanceof DefaultTrigger)) {
throw new UnsupportedOperationException(
- "Stateful ParDo over merging windows" + NOT_SUPPORTED);
+ "Custom triggers (" + strategy.getTrigger() + ")" + NOT_SUPPORTED);
+ }
+ if (strategy.getMode() != AccumulationMode.DISCARDING_FIRED_PANES) {
+ throw new UnsupportedOperationException("Accumulating fired panes" +
NOT_SUPPORTED);
+ }
+ }
+
+ private static void checkNotMergingWindows(
+ WindowingStrategy<?, ?> windowingStrategy, String feature) {
+ if (windowingStrategy.needsMerge()) {
+ throw new UnsupportedOperationException(feature + " over merging
windows" + NOT_SUPPORTED);
}
}
diff --git
a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTranslator.java
b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/GroupByKeyStreamingTranslator.java
similarity index 57%
copy from
runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTranslator.java
copy to
runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/GroupByKeyStreamingTranslator.java
index 199570ae64a..b997f0d7498 100644
---
a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTranslator.java
+++
b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/GroupByKeyStreamingTranslator.java
@@ -18,20 +18,14 @@
package
org.apache.beam.runners.spark.structuredstreaming.translation.streaming;
import org.apache.beam.runners.spark.coders.CoderHelpers;
-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.DoFnRunnerFactory;
-import
org.apache.beam.runners.spark.structuredstreaming.translation.batch.functions.SparkSideInputReader;
-import
org.apache.beam.runners.spark.structuredstreaming.translation.streaming.state.BeamStatefulProcessor;
+import
org.apache.beam.runners.spark.structuredstreaming.translation.streaming.state.GroupByKeyStatefulProcessor;
import org.apache.beam.sdk.annotations.Internal;
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.GroupByKey;
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.spark.api.java.function.MapFunction;
@@ -40,41 +34,31 @@ import org.apache.spark.sql.Encoders;
import org.apache.spark.sql.streaming.OutputMode;
import org.apache.spark.sql.streaming.TimeMode;
-/** Translates a stateful {@link ParDo.MultiOutput} for Spark 4 Structured
Streaming. */
+/** Translates {@link GroupByKey} for Spark 4 Structured Streaming via
transformWithState. */
@Internal
-public class StatefulParDoStreamingTranslator<K, V, OutputT>
+public class GroupByKeyStreamingTranslator<K, V>
extends TransformTranslator<
- PCollection<? extends KV<K, V>>, PCollectionTuple,
ParDo.MultiOutput<KV<K, V>, OutputT>> {
+ PCollection<KV<K, V>>, PCollection<KV<K, Iterable<V>>>, GroupByKey<K,
V>> {
- public StatefulParDoStreamingTranslator() {
+ public GroupByKeyStreamingTranslator() {
super(0.2f);
}
@Override
- protected void translate(ParDo.MultiOutput<KV<K, V>, OutputT> transform,
Context cxt) {
- @SuppressWarnings("unchecked")
- PCollection<KV<K, V>> input = (PCollection<KV<K, V>>) cxt.getInput();
+ protected void translate(GroupByKey<K, V> transform, Context cxt) {
+ PCollection<KV<K, V>> input = cxt.getInput();
WindowingStrategy<?, ?> windowing = input.getWindowingStrategy();
- DoFn<KV<K, V>, OutputT> doFn = transform.getFn();
-
KvCoder<K, V> inputCoder = (KvCoder<K, V>) input.getCoder();
Coder<K> keyCoder = inputCoder.getKeyCoder();
- TupleTag<OutputT> mainOutputTag = transform.getMainOutputTag();
- PCollection<OutputT> output = cxt.getOutput(mainOutputTag);
- Coder<OutputT> outputCoder = output.getCoder();
-
- DoFnRunnerFactory<KV<K, V>, OutputT> runnerFactory =
- DoFnRunnerFactory.simple(
- cxt.getCurrentTransform(), input, SparkSideInputReader.empty(),
false);
- MetricsAccumulator metrics =
MetricsAccumulator.getInstance(cxt.getSparkSession());
+ PCollection<KV<K, Iterable<V>>> output = cxt.getOutput();
+ Coder<KV<K, Iterable<V>>> outputCoder = output.getCoder();
- BeamStatefulProcessor<K, V, OutputT> processor =
- new BeamStatefulProcessor<>(
- doFn, runnerFactory, inputCoder, windowing,
cxt.getOptionsSupplier(), metrics);
+ GroupByKeyStatefulProcessor<K, V> processor =
+ new GroupByKeyStatefulProcessor<>(inputCoder, windowing,
cxt.getOptionsSupplier());
MapFunction<WindowedValue<KV<K, V>>, byte[]> keyFn =
v -> CoderHelpers.toByteArray(v.getValue().getKey(), keyCoder);
- Dataset<WindowedValue<OutputT>> result =
+ Dataset<WindowedValue<KV<K, Iterable<V>>>> result =
cxt.getDataset(input)
.groupByKey(keyFn, Encoders.BINARY())
.transformWithState(
diff --git
a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTranslator.java
b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTranslator.java
index 199570ae64a..3f43a6b25e6 100644
---
a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTranslator.java
+++
b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTranslator.java
@@ -22,7 +22,7 @@ import
org.apache.beam.runners.spark.structuredstreaming.metrics.MetricsAccumula
import
org.apache.beam.runners.spark.structuredstreaming.translation.TransformTranslator;
import
org.apache.beam.runners.spark.structuredstreaming.translation.batch.DoFnRunnerFactory;
import
org.apache.beam.runners.spark.structuredstreaming.translation.batch.functions.SparkSideInputReader;
-import
org.apache.beam.runners.spark.structuredstreaming.translation.streaming.state.BeamStatefulProcessor;
+import
org.apache.beam.runners.spark.structuredstreaming.translation.streaming.state.ParDoStatefulProcessor;
import org.apache.beam.sdk.annotations.Internal;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.KvCoder;
@@ -68,8 +68,8 @@ public class StatefulParDoStreamingTranslator<K, V, OutputT>
cxt.getCurrentTransform(), input, SparkSideInputReader.empty(),
false);
MetricsAccumulator metrics =
MetricsAccumulator.getInstance(cxt.getSparkSession());
- BeamStatefulProcessor<K, V, OutputT> processor =
- new BeamStatefulProcessor<>(
+ ParDoStatefulProcessor<K, V, OutputT> processor =
+ new ParDoStatefulProcessor<>(
doFn, runnerFactory, inputCoder, windowing,
cxt.getOptionsSupplier(), metrics);
MapFunction<WindowedValue<KV<K, V>>, byte[]> keyFn =
diff --git
a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessor.java
b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessor.java
index 7d794ef0450..568c41d23ec 100644
---
a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessor.java
+++
b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessor.java
@@ -25,29 +25,19 @@ import java.util.Comparator;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
-import java.util.function.Supplier;
-import org.apache.beam.runners.core.DoFnRunner;
-import org.apache.beam.runners.core.DoFnRunners;
-import org.apache.beam.runners.core.StatefulDoFnRunner;
+import java.util.stream.Collectors;
import org.apache.beam.runners.core.TimerInternals.TimerData;
import org.apache.beam.runners.core.TimerInternals.TimerDataCoderV2;
import org.apache.beam.runners.spark.coders.CoderHelpers;
import org.apache.beam.runners.spark.stateful.SparkStateInternals;
import org.apache.beam.runners.spark.stateful.SparkStateInternals.StateCells;
import org.apache.beam.runners.spark.stateful.SparkTimerInternals;
-import
org.apache.beam.runners.spark.structuredstreaming.metrics.MetricsAccumulator;
-import
org.apache.beam.runners.spark.structuredstreaming.translation.batch.DoFnRunnerFactory;
-import
org.apache.beam.runners.spark.structuredstreaming.translation.batch.DoFnRunnerFactory.DoFnRunnerWithTeardown;
-import
org.apache.beam.runners.spark.structuredstreaming.translation.batch.StatefulTaskRunner;
import
org.apache.beam.runners.spark.structuredstreaming.translation.utils.ScalaInterop;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.coders.ListCoder;
-import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.state.TimeDomain;
-import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
-import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
import org.apache.beam.sdk.util.WindowedValueMultiReceiver;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.TupleTag;
@@ -67,43 +57,45 @@ import org.checkerframework.checker.nullness.qual.Nullable;
import org.joda.time.Instant;
import scala.collection.Iterator;
-/** Runs a stateful {@link DoFn} within Spark 4 Structured Streaming {@code
transformWithState}. */
-public class BeamStatefulProcessor<K, V, OutputT>
+/** Base stateful processor bridging Beam state and timers over Spark
transformWithState. */
+public abstract class BeamStatefulProcessor<K, V, OutputT>
extends StatefulProcessor<byte[], WindowedValue<KV<K, V>>,
WindowedValue<OutputT>> {
private static final String BEAM_STATE = "beamState";
private static final String BEAM_TIMERS = "beamTimers";
private static final String BEAM_WATERMARK = "beamWatermark";
- private final DoFn<KV<K, V>, OutputT> doFn;
- private final DoFnRunnerFactory<KV<K, V>, OutputT> runnerFactory;
- private final KvCoder<K, V> inputCoder;
- private final WindowingStrategy<?, ?> windowingStrategy;
- private final Supplier<PipelineOptions> optionsSupplier;
- private final MetricsAccumulator metrics;
- private final Coder<BoundedWindow> windowCoder;
+ protected final KvCoder<K, V> inputCoder;
+ protected final WindowingStrategy<?, ?> windowingStrategy;
+ protected final Coder<BoundedWindow> windowCoder;
private final TimerDataCoderV2 timerDataCoder;
- private final StatefulTaskRunner<KV<K, V>, OutputT> taskRunner = new
StatefulTaskRunner<>();
private transient @Nullable MapState<String, byte[]> beamState;
private transient @Nullable ValueState<byte[]> beamTimers;
private transient @Nullable ValueState<Long> beamWatermark;
private transient @Nullable List<WindowedValue<OutputT>> currentOutputs;
- private transient boolean needsBundleStart;
+
+ private transient @Nullable WindowedValueMultiReceiver outputReceiver;
+
+ protected WindowedValueMultiReceiver outputReceiver() {
+ if (outputReceiver == null) {
+ outputReceiver =
+ new WindowedValueMultiReceiver() {
+ @Override
+ public <T> void output(TupleTag<T> tag, WindowedValue<T> output) {
+ @SuppressWarnings("unchecked")
+ WindowedValue<OutputT> out = (WindowedValue<OutputT>) output;
+ checkStateNotNull(currentOutputs, "currentOutputs not
initialized").add(out);
+ }
+ };
+ }
+ return outputReceiver;
+ }
public BeamStatefulProcessor(
- DoFn<KV<K, V>, OutputT> doFn,
- DoFnRunnerFactory<KV<K, V>, OutputT> runnerFactory,
- KvCoder<K, V> inputCoder,
- WindowingStrategy<?, ?> windowingStrategy,
- Supplier<PipelineOptions> optionsSupplier,
- MetricsAccumulator metrics) {
- this.doFn = doFn;
- this.runnerFactory = runnerFactory;
+ KvCoder<K, V> inputCoder, WindowingStrategy<?, ?> windowingStrategy) {
this.inputCoder = inputCoder;
this.windowingStrategy = windowingStrategy;
- this.optionsSupplier = optionsSupplier;
- this.metrics = metrics;
@SuppressWarnings("unchecked")
Coder<BoundedWindow> windowCoder =
(Coder<BoundedWindow>) windowingStrategy.getWindowFn().windowCoder();
@@ -151,16 +143,13 @@ public class BeamStatefulProcessor<K, V, OutputT>
byte[] keyBytes, Iterator<WindowedValue<KV<K, V>>> rows, TimerValues
timerValues) {
K key = CoderHelpers.fromByteArray(keyBytes, inputCoder.getKeyCoder());
List<WindowedValue<OutputT>> outputs = new ArrayList<>();
+ this.currentOutputs = outputs;
SparkTimerInternals timerInternals = timerInternals(timerValues);
+ SparkStateInternals<K> state = stateInternals(key);
- DoFnRunner<KV<K, V>, OutputT> runner = keyRunner(key, timerInternals,
outputs);
- while (rows.hasNext()) {
- runner.processElement(rows.next());
- }
+ processElements(key, rows, timerInternals, state);
timerInternals.advanceWatermark();
- // The next bundle opens in keyRunner.
- runner.finishBundle();
- needsBundleStart = true;
+ processTimers(key, timerInternals, state);
persistTimers(timerInternals);
reconcileWakeupTimer(timerInternals, null);
@@ -173,14 +162,12 @@ public class BeamStatefulProcessor<K, V, OutputT>
byte[] keyBytes, TimerValues timerValues, ExpiredTimerInfo
expiredTimerInfo) {
K key = CoderHelpers.fromByteArray(keyBytes, inputCoder.getKeyCoder());
List<WindowedValue<OutputT>> outputs = new ArrayList<>();
+ this.currentOutputs = outputs;
SparkTimerInternals timerInternals = timerInternals(timerValues);
+ SparkStateInternals<K> state = stateInternals(key);
timerInternals.advanceWatermark();
- DoFnRunner<KV<K, V>, OutputT> runner = keyRunner(key, timerInternals,
outputs);
- fireDueTimers(key, timerInternals.currentInputWatermarkTime(),
timerInternals, runner);
- // The next bundle opens in keyRunner.
- runner.finishBundle();
- needsBundleStart = true;
+ processTimers(key, timerInternals, state);
persistTimers(timerInternals);
reconcileWakeupTimer(timerInternals, expiredTimerInfo.getExpiryTimeInMs());
@@ -188,70 +175,28 @@ public class BeamStatefulProcessor<K, V, OutputT>
return ScalaInterop.scalaIterator(outputs);
}
- @Override
- public void close() {
- taskRunner.teardownOnce();
+ private SparkStateInternals<K> stateInternals(K key) {
+ return SparkStateInternals.forKey(
+ key, new MapStateAdapter(checkStateNotNull(beamState, "beamState not
initialized")));
}
- /** Creates the task runner on first use, {@code create} opens the first
bundle. */
- private DoFnRunnerWithTeardown<KV<K, V>, OutputT> baseRunner(
- SparkStateInternals<K> stateInternals, SparkTimerInternals
timerInternals) {
- return taskRunner.getOrCreate(
- ctx -> {
- ctx.set(stateInternals, timerInternals);
- return runnerFactory.create(
- optionsSupplier.get(),
- metrics,
- new WindowedValueMultiReceiver() {
- @Override
- public <T> void output(TupleTag<T> tag, WindowedValue<T>
output) {
- @SuppressWarnings("unchecked")
- WindowedValue<OutputT> out = (WindowedValue<OutputT>) output;
- checkStateNotNull(currentOutputs, "currentOutputs not
initialized").add(out);
- }
- },
- ctx);
- });
+ protected List<TimerData> dueEventTimeTimers(SparkTimerInternals
timerInternals) {
+ Instant watermark = timerInternals.currentInputWatermarkTime();
+ return timerInternals.getTimers().stream()
+ .filter(
+ t -> t.getDomain().equals(TimeDomain.EVENT_TIME) &&
watermark.isAfter(t.getTimestamp()))
+ .sorted(Comparator.comparing(TimerData::getTimestamp))
+ .collect(Collectors.toList());
}
- private DoFnRunner<KV<K, V>, OutputT> keyRunner(
- K key, SparkTimerInternals timerInternals, List<WindowedValue<OutputT>>
outputs) {
- this.currentOutputs = outputs;
- MapState<String, byte[]> state = checkStateNotNull(beamState);
- SparkStateInternals<K> stateInternals =
- SparkStateInternals.forKey(key, new MapStateAdapter(state));
- DoFnRunnerWithTeardown<KV<K, V>, OutputT> base =
baseRunner(stateInternals, timerInternals);
- taskRunner.stepContext().set(stateInternals, timerInternals);
-
- // Closed in handleInputRows and handleExpiredTimer.
- if (needsBundleStart) {
- needsBundleStart = false;
- base.startBundle();
- }
+ protected abstract void processElements(
+ K key,
+ Iterator<WindowedValue<KV<K, V>>> rows,
+ SparkTimerInternals timers,
+ SparkStateInternals<K> state);
- StatefulDoFnRunner.CleanupTimer<KV<K, V>> cleanupTimer =
- new StatefulDoFnRunner.TimeInternalsCleanupTimer<KV<K, V>>(
- timerInternals, windowingStrategy) {
- @Override
- public void setForWindow(KV<K, V> input, BoundedWindow window) {
- // GlobalWindow state is never garbage collected, as in the Flink
runner.
- if (!window.equals(GlobalWindow.INSTANCE)) {
- super.setForWindow(input, window);
- }
- }
- };
- StatefulDoFnRunner.StateCleaner<BoundedWindow> stateCleaner =
- new StatefulDoFnRunner.StateInternalsStateCleaner<>(doFn,
stateInternals, windowCoder);
-
- return DoFnRunners.defaultStatefulDoFnRunner(
- doFn,
- inputCoder,
- base,
- taskRunner.stepContext(),
- windowingStrategy,
- cleanupTimer,
- stateCleaner);
- }
+ protected abstract void processTimers(
+ K key, SparkTimerInternals timers, SparkStateInternals<K> state);
private void restoreTimers(SparkTimerInternals timerInternals) {
ValueState<byte[]> timersState = checkStateNotNull(beamTimers);
@@ -278,28 +223,6 @@ public class BeamStatefulProcessor<K, V, OutputT>
}
}
- private void fireDueTimers(
- K key,
- Instant watermark,
- SparkTimerInternals timerInternals,
- DoFnRunner<KV<K, V>, OutputT> runner) {
- while (true) {
- TimerData next =
- timerInternals.getTimers().stream()
- .filter(
- t ->
- t.getDomain().equals(TimeDomain.EVENT_TIME)
- && watermark.isAfter(t.getTimestamp()))
- .min(Comparator.comparing(TimerData::getTimestamp))
- .orElse(null);
- if (next == null) {
- break;
- }
- timerInternals.deleteTimer(next);
- StatefulTaskRunner.fireTimer(runner, key, next);
- }
- }
-
private void reconcileWakeupTimer(
SparkTimerInternals timerInternals, @Nullable Long firedExpiryMs) {
Long nextWakeupMs = null;
diff --git
a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/GroupByKeyStatefulProcessor.java
b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/GroupByKeyStatefulProcessor.java
new file mode 100644
index 00000000000..cbe592af0ea
--- /dev/null
+++
b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/GroupByKeyStatefulProcessor.java
@@ -0,0 +1,141 @@
+/*
+ * 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.streaming.state;
+
+import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.function.Supplier;
+import org.apache.beam.runners.core.DoFnRunner;
+import org.apache.beam.runners.core.DoFnRunners;
+import org.apache.beam.runners.core.GroupAlsoByWindowViaWindowSetNewDoFn;
+import org.apache.beam.runners.core.KeyedWorkItem;
+import org.apache.beam.runners.core.KeyedWorkItemCoder;
+import org.apache.beam.runners.core.KeyedWorkItems;
+import org.apache.beam.runners.core.NullSideInputReader;
+import org.apache.beam.runners.core.SystemReduceFn;
+import org.apache.beam.runners.core.TimerInternals.TimerData;
+import org.apache.beam.runners.spark.stateful.SparkStateInternals;
+import org.apache.beam.runners.spark.stateful.SparkTimerInternals;
+import
org.apache.beam.runners.spark.structuredstreaming.translation.batch.StatefulTaskRunner.MutableStepContext;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFnSchemaInformation;
+import org.apache.beam.sdk.util.WindowedValueMultiReceiver;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.WindowedValue;
+import org.apache.beam.sdk.values.WindowedValues;
+import org.apache.beam.sdk.values.WindowingStrategy;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import scala.collection.Iterator;
+
+/** Runs a {@link GroupAlsoByWindowViaWindowSetNewDoFn} in {@code
transformWithState}. */
+public class GroupByKeyStatefulProcessor<K, V>
+ extends BeamStatefulProcessor<K, V, KV<K, Iterable<V>>> {
+
+ private final Supplier<PipelineOptions> optionsSupplier;
+ private final TupleTag<KV<K, Iterable<V>>> mainTag = new
TupleTag<>("output");
+
+ private transient @Nullable DoFnRunner<KeyedWorkItem<K, V>, KV<K,
Iterable<V>>> runner;
+ private transient @Nullable MutableStepContext stepContext;
+
+ public GroupByKeyStatefulProcessor(
+ KvCoder<K, V> inputCoder,
+ WindowingStrategy<?, ?> windowingStrategy,
+ Supplier<PipelineOptions> optionsSupplier) {
+ super(inputCoder, windowingStrategy);
+ this.optionsSupplier = optionsSupplier;
+ }
+
+ @Override
+ protected void processElements(
+ K key,
+ Iterator<WindowedValue<KV<K, V>>> rows,
+ SparkTimerInternals timerInternals,
+ SparkStateInternals<K> stateInternals) {
+ DoFnRunner<KeyedWorkItem<K, V>, KV<K, Iterable<V>>> r = runner();
+ checkStateNotNull(stepContext).set(stateInternals, timerInternals);
+
+ List<WindowedValue<V>> elementValues = new ArrayList<>();
+ while (rows.hasNext()) {
+ WindowedValue<KV<K, V>> row = rows.next();
+ elementValues.add(row.withValue(row.getValue().getValue()));
+ }
+ r.processElement(
+
WindowedValues.valueInGlobalWindow(KeyedWorkItems.elementsWorkItem(key,
elementValues)));
+ }
+
+ @Override
+ protected void processTimers(
+ K key, SparkTimerInternals timerInternals, SparkStateInternals<K>
stateInternals) {
+ List<TimerData> dueTimers = dueEventTimeTimers(timerInternals);
+ if (dueTimers.isEmpty()) {
+ return;
+ }
+ for (TimerData timer : dueTimers) {
+ timerInternals.deleteTimer(timer);
+ }
+ DoFnRunner<KeyedWorkItem<K, V>, KV<K, Iterable<V>>> r = runner();
+ checkStateNotNull(stepContext).set(stateInternals, timerInternals);
+ r.processElement(
+ WindowedValues.valueInGlobalWindow(KeyedWorkItems.timersWorkItem(key,
dueTimers)));
+ }
+
+ private DoFnRunner<KeyedWorkItem<K, V>, KV<K, Iterable<V>>> runner() {
+ if (runner == null) {
+ MutableStepContext ctx = new MutableStepContext();
+ this.stepContext = ctx;
+ WindowedValueMultiReceiver receiver = outputReceiver();
+ DoFn<KeyedWorkItem<K, V>, KV<K, Iterable<V>>> gabwDoFn =
+ GroupAlsoByWindowViaWindowSetNewDoFn.create(
+ windowingStrategy,
+ k -> ctx.stateInternals(),
+ k -> ctx.timerInternals(),
+ NullSideInputReader.empty(),
+ SystemReduceFn.buffering(inputCoder.getValueCoder()),
+ receiver,
+ mainTag);
+
+ Coder<KeyedWorkItem<K, V>> itemCoder =
+ KeyedWorkItemCoder.of(inputCoder.getKeyCoder(),
inputCoder.getValueCoder(), windowCoder);
+
+ DoFnRunner<KeyedWorkItem<K, V>, KV<K, Iterable<V>>> simpleRunner =
+ DoFnRunners.simpleRunner(
+ optionsSupplier.get(),
+ gabwDoFn,
+ NullSideInputReader.empty(),
+ receiver,
+ mainTag,
+ Collections.emptyList(),
+ ctx,
+ itemCoder,
+ Collections.emptyMap(),
+ windowingStrategy,
+ DoFnSchemaInformation.create(),
+ Collections.emptyMap());
+
+ runner = DoFnRunners.lateDataDroppingRunner(simpleRunner, ctx,
windowingStrategy);
+ }
+ return runner;
+ }
+}
diff --git
a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/ParDoStatefulProcessor.java
b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/ParDoStatefulProcessor.java
new file mode 100644
index 00000000000..f25c523eaac
--- /dev/null
+++
b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/ParDoStatefulProcessor.java
@@ -0,0 +1,156 @@
+/*
+ * 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.streaming.state;
+
+import java.util.List;
+import java.util.function.Supplier;
+import org.apache.beam.runners.core.DoFnRunner;
+import org.apache.beam.runners.core.DoFnRunners;
+import org.apache.beam.runners.core.StatefulDoFnRunner;
+import org.apache.beam.runners.core.TimerInternals.TimerData;
+import org.apache.beam.runners.spark.stateful.SparkStateInternals;
+import org.apache.beam.runners.spark.stateful.SparkTimerInternals;
+import
org.apache.beam.runners.spark.structuredstreaming.metrics.MetricsAccumulator;
+import
org.apache.beam.runners.spark.structuredstreaming.translation.batch.DoFnRunnerFactory;
+import
org.apache.beam.runners.spark.structuredstreaming.translation.batch.DoFnRunnerFactory.DoFnRunnerWithTeardown;
+import
org.apache.beam.runners.spark.structuredstreaming.translation.batch.StatefulTaskRunner;
+import org.apache.beam.sdk.coders.KvCoder;
+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.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.WindowedValue;
+import org.apache.beam.sdk.values.WindowingStrategy;
+import scala.collection.Iterator;
+
+/** Runs a stateful {@link DoFn} within Spark 4 Structured Streaming {@code
transformWithState}. */
+public class ParDoStatefulProcessor<K, V, OutputT> extends
BeamStatefulProcessor<K, V, OutputT> {
+
+ private final DoFn<KV<K, V>, OutputT> doFn;
+ private final DoFnRunnerFactory<KV<K, V>, OutputT> runnerFactory;
+ private final Supplier<PipelineOptions> optionsSupplier;
+ private final MetricsAccumulator metrics;
+ private final StatefulTaskRunner<KV<K, V>, OutputT> taskRunner = new
StatefulTaskRunner<>();
+
+ private transient boolean needsBundleStart;
+
+ public ParDoStatefulProcessor(
+ DoFn<KV<K, V>, OutputT> doFn,
+ DoFnRunnerFactory<KV<K, V>, OutputT> runnerFactory,
+ KvCoder<K, V> inputCoder,
+ WindowingStrategy<?, ?> windowingStrategy,
+ Supplier<PipelineOptions> optionsSupplier,
+ MetricsAccumulator metrics) {
+ super(inputCoder, windowingStrategy);
+ this.doFn = doFn;
+ this.runnerFactory = runnerFactory;
+ this.optionsSupplier = optionsSupplier;
+ this.metrics = metrics;
+ }
+
+ @Override
+ protected void processElements(
+ K key,
+ Iterator<WindowedValue<KV<K, V>>> rows,
+ SparkTimerInternals timerInternals,
+ SparkStateInternals<K> stateInternals) {
+ DoFnRunner<KV<K, V>, OutputT> runner = keyRunner(stateInternals,
timerInternals);
+ while (rows.hasNext()) {
+ runner.processElement(rows.next());
+ }
+ // The next bundle opens in keyRunner.
+ runner.finishBundle();
+ needsBundleStart = true;
+ }
+
+ @Override
+ protected void processTimers(
+ K key, SparkTimerInternals timerInternals, SparkStateInternals<K>
stateInternals) {
+ if (dueEventTimeTimers(timerInternals).isEmpty()) {
+ return;
+ }
+ DoFnRunner<KV<K, V>, OutputT> runner = keyRunner(stateInternals,
timerInternals);
+ fireDueTimers(key, timerInternals, runner);
+ // The next bundle opens in keyRunner.
+ runner.finishBundle();
+ needsBundleStart = true;
+ }
+
+ @Override
+ public void close() {
+ taskRunner.teardownOnce();
+ }
+
+ private void fireDueTimers(
+ K key, SparkTimerInternals timerInternals, DoFnRunner<KV<K, V>, OutputT>
runner) {
+ while (true) {
+ List<TimerData> due = dueEventTimeTimers(timerInternals);
+ if (due.isEmpty()) {
+ break;
+ }
+ TimerData next = due.get(0);
+ timerInternals.deleteTimer(next);
+ StatefulTaskRunner.fireTimer(runner, key, next);
+ }
+ }
+
+ private DoFnRunnerWithTeardown<KV<K, V>, OutputT> baseRunner(
+ SparkStateInternals<K> stateInternals, SparkTimerInternals
timerInternals) {
+ return taskRunner.getOrCreate(
+ ctx -> {
+ ctx.set(stateInternals, timerInternals);
+ return runnerFactory.create(optionsSupplier.get(), metrics,
outputReceiver(), ctx);
+ });
+ }
+
+ private DoFnRunner<KV<K, V>, OutputT> keyRunner(
+ SparkStateInternals<K> stateInternals, SparkTimerInternals
timerInternals) {
+ DoFnRunnerWithTeardown<KV<K, V>, OutputT> base =
baseRunner(stateInternals, timerInternals);
+ taskRunner.stepContext().set(stateInternals, timerInternals);
+
+ // Closed in processElements and processTimers.
+ if (needsBundleStart) {
+ needsBundleStart = false;
+ base.startBundle();
+ }
+
+ StatefulDoFnRunner.CleanupTimer<KV<K, V>> cleanupTimer =
+ new StatefulDoFnRunner.TimeInternalsCleanupTimer<KV<K, V>>(
+ timerInternals, windowingStrategy) {
+ @Override
+ public void setForWindow(KV<K, V> input, BoundedWindow window) {
+ // GlobalWindow state is never garbage collected, as in the Flink
runner.
+ if (!window.equals(GlobalWindow.INSTANCE)) {
+ super.setForWindow(input, window);
+ }
+ }
+ };
+ StatefulDoFnRunner.StateCleaner<BoundedWindow> stateCleaner =
+ new StatefulDoFnRunner.StateInternalsStateCleaner<>(doFn,
stateInternals, windowCoder);
+
+ return DoFnRunners.defaultStatefulDoFnRunner(
+ doFn,
+ inputCoder,
+ base,
+ taskRunner.stepContext(),
+ windowingStrategy,
+ cleanupTimer,
+ stateCleaner);
+ }
+}
diff --git
a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreamingTest.java
b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreamingTest.java
index 39992b76fe3..a5afa7c1a4a 100644
---
a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreamingTest.java
+++
b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreamingTest.java
@@ -33,13 +33,16 @@ import org.apache.beam.sdk.state.StateSpecs;
import org.apache.beam.sdk.state.TimeDomain;
import org.apache.beam.sdk.state.TimerSpec;
import org.apache.beam.sdk.state.TimerSpecs;
-import org.apache.beam.sdk.transforms.Combine;
+import org.apache.beam.sdk.transforms.Count;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.GroupByKey;
import org.apache.beam.sdk.transforms.ParDo;
-import org.apache.beam.sdk.transforms.Sum;
+import org.apache.beam.sdk.transforms.windowing.AfterPane;
+import org.apache.beam.sdk.transforms.windowing.AfterWatermark;
import org.apache.beam.sdk.transforms.windowing.FixedWindows;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindows;
+import org.apache.beam.sdk.transforms.windowing.Repeatedly;
import org.apache.beam.sdk.transforms.windowing.Sessions;
import org.apache.beam.sdk.transforms.windowing.Window;
import org.apache.beam.sdk.values.KV;
@@ -84,17 +87,6 @@ public class PipelineTranslatorStreamingTest implements
Serializable {
throw new AssertionError("missing " + expected);
}
- @Test
- public void rejectsGroupByKey() throws Exception {
- assertUnsupported(kv("g").apply(GroupByKey.create()).getPipeline(),
"GroupByKey");
- }
-
- @Test
- public void rejectsCombinePerKey() throws Exception {
- assertUnsupported(
- kv("c").apply(Combine.perKey(Sum.ofIntegers())).getPipeline(),
"Combine.perKey");
- }
-
@Test
public void rejectsImpulseFromCreate() throws Exception {
SparkStructuredStreamingPipelineOptions o =
StreamingTestUtils.streamingOptions(temp);
@@ -136,6 +128,43 @@ public class PipelineTranslatorStreamingTest implements
Serializable {
.apply(ParDo.of(new StatefulDoFn()))
.getPipeline(),
"merging windows"
+ },
+ {
+ kv("gbkmw")
+ .apply(Window.into(Sessions.withGapDuration(Duration.millis(10))))
+ .apply(GroupByKey.create())
+ .getPipeline(),
+ "merging windows"
+ },
+ {
+ kv("trig")
+ .apply(
+ Window.<KV<String,
Integer>>into(FixedWindows.of(Duration.millis(1)))
+
.triggering(Repeatedly.forever(AfterWatermark.pastEndOfWindow()))
+ .withAllowedLateness(Duration.millis(1))
+ .discardingFiredPanes())
+ .apply(GroupByKey.create())
+ .getPipeline(),
+ "Custom triggers"
+ },
+ {
+ kv("accum")
+ .apply(
+ Window.<KV<String,
Integer>>into(FixedWindows.of(Duration.millis(1)))
+ .accumulatingFiredPanes())
+ .apply(GroupByKey.create())
+ .getPipeline(),
+ "Accumulating fired panes"
+ },
+ {
+ kv("cgdef")
+ .apply(
+ Window.<KV<String, Integer>>into(new GlobalWindows())
+
.triggering(Repeatedly.forever(AfterPane.elementCountAtLeast(1)))
+ .discardingFiredPanes())
+ .apply(Count.globally())
+ .getPipeline(),
+ "Combine.globally() with a default value (use withoutDefaults())"
}
};
diff --git
a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTest.java
b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTest.java
index 18143ce4d92..5679950e189 100644
---
a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTest.java
+++
b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTest.java
@@ -21,7 +21,6 @@ import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;
-import java.io.File;
import java.io.Serializable;
import java.util.Arrays;
import java.util.Set;
@@ -213,22 +212,10 @@ public class StatefulParDoStreamingTest implements
Serializable {
Set<String> collected = StreamingTestUtils.collected(collectorId);
assertTrue(collected.contains(TestUnboundedSource.element(TAG_EXPIRY, 0,
0)));
for (String stateVarName : Arrays.asList("beamState", "beamTimers",
"beamWatermark")) {
- assertEquals(stateVarName, 0L, stateRows(checkpoint, stateVarName));
+ assertEquals(stateVarName, 0L, StreamingTestUtils.stateRows(checkpoint,
stateVarName));
}
}
- /** Rows of a state variable of the query of the only leaf, checkpointed
under 0. */
- private static long stateRows(String checkpoint, String stateVarName) {
- return SESSION
- .getSession()
- .read()
- .format("statestore")
- .option("path", new File(checkpoint, "0").getAbsolutePath())
- .option("stateVarName", stateVarName)
- .load()
- .count();
- }
-
private static final class KeySelectorFn extends DoFn<String, KV<String,
String>> {
@ProcessElement
public void process(@Element String element, OutputReceiver<KV<String,
String>> out) {
diff --git
a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTestUtils.java
b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTestUtils.java
index 3d784115415..5616b5b4b7c 100644
---
a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTestUtils.java
+++
b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTestUtils.java
@@ -17,6 +17,7 @@
*/
package
org.apache.beam.runners.spark.structuredstreaming.translation.streaming;
+import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
@@ -33,6 +34,7 @@ import org.apache.beam.sdk.PipelineResult;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.transforms.DoFn;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
+import org.apache.spark.sql.SparkSession;
import org.joda.time.Duration;
import org.junit.rules.TemporaryFolder;
@@ -153,4 +155,15 @@ public final class StreamingTestUtils {
options.setCheckpointDir(checkpointPath);
return options;
}
+
+ /** Rows of a state variable of the query of the only leaf, checkpointed
under 0. */
+ static long stateRows(String checkpoint, String stateVarName) {
+ return SparkSession.active()
+ .read()
+ .format("statestore")
+ .option("path", new File(checkpoint, "0").getAbsolutePath())
+ .option("stateVarName", stateVarName)
+ .load()
+ .count();
+ }
}
diff --git
a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/WindowedGroupByKeyStreamingTest.java
b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/WindowedGroupByKeyStreamingTest.java
new file mode 100644
index 00000000000..2b0305fa66d
--- /dev/null
+++
b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/WindowedGroupByKeyStreamingTest.java
@@ -0,0 +1,274 @@
+/*
+ * 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.streaming;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+import java.io.Serializable;
+import java.util.Arrays;
+import java.util.Set;
+import org.apache.beam.runners.spark.StreamingTest;
+import org.apache.beam.runners.spark.structuredstreaming.SparkSessionRule;
+import
org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions;
+import
org.apache.beam.runners.spark.structuredstreaming.io.streaming.TestUnboundedSource;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.PipelineResult;
+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.io.Read;
+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.transforms.Combine;
+import org.apache.beam.sdk.transforms.Count;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.GroupByKey;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.windowing.FixedWindows;
+import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
+import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.transforms.windowing.SlidingWindows;
+import org.apache.beam.sdk.transforms.windowing.Window;
+import org.apache.beam.sdk.values.KV;
+import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.junit.After;
+import org.junit.ClassRule;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+import org.junit.rules.TemporaryFolder;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** End to end tests for windowed GroupByKey, Combine.perKey and
Combine.globally on Spark 4. */
+@RunWith(JUnit4.class)
+@Category(StreamingTest.class)
+public class WindowedGroupByKeyStreamingTest implements Serializable {
+
+ @ClassRule public static final SparkSessionRule SESSION = new
SparkSessionRule();
+
+ @Rule public transient TemporaryFolder tempFolder = new TemporaryFolder();
+
+ private static final String TAG_FIXED = "fixed-gbk";
+ private static final String TAG_SLIDING = "sliding-combine";
+ private static final String TAG_LATE = "late-dropped";
+ private static final String TAG_CHAINED = "chained-timer";
+ private static final String TAG_GLOBALLY = "globally-combine";
+ private static final String TAG_EXPIRY = "gbk-expiry";
+
+ @After
+ public void tearDown() {
+ TestUnboundedSource.forget(TAG_FIXED);
+ TestUnboundedSource.forget(TAG_SLIDING);
+ TestUnboundedSource.forget(TAG_LATE);
+ TestUnboundedSource.forget(TAG_CHAINED);
+ TestUnboundedSource.forget(TAG_GLOBALLY);
+ TestUnboundedSource.forget(TAG_EXPIRY);
+ }
+
+ private SparkStructuredStreamingPipelineOptions options(String tag) throws
Exception {
+ SparkStructuredStreamingPipelineOptions options =
+
StreamingTestUtils.streamingOptions(tempFolder.newFolder(tag).getAbsolutePath());
+ SESSION.configure(options);
+ options.setMaxRecordsPerBatch(1L);
+ return options;
+ }
+
+ @Test
+ public void testFixedWindowsGroupByKey() throws Exception {
+ String col = StreamingTestUtils.newCollectorId(TAG_FIXED);
+ Pipeline p = Pipeline.create(options(TAG_FIXED));
+ p.apply(Read.from(new TestUnboundedSource(TAG_FIXED, 1, 6)))
+ .apply(ParDo.of(new ToKvFn()))
+ .setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()))
+ .apply(Window.into(FixedWindows.of(Duration.standardSeconds(2))))
+ .apply(GroupByKey.create())
+ .apply(ParDo.of(new FormatGbkFn()))
+ .apply(ParDo.of(new StreamingTestUtils.CollectDoFn<>(col)));
+
+ assertEquals(PipelineResult.State.DONE,
StreamingTestUtils.run(p).getState());
+ Set<String> res = StreamingTestUtils.collected(col);
+ assertTrue(res.contains(TestUnboundedSource.BASE_MILLIS + ":k=2"));
+ assertTrue(res.contains(TestUnboundedSource.BASE_MILLIS + 2000 + ":k=2"));
+ }
+
+ @Test
+ public void testSlidingWindowsCombinePerKey() throws Exception {
+ String col = StreamingTestUtils.newCollectorId(TAG_SLIDING);
+ Pipeline p = Pipeline.create(options(TAG_SLIDING));
+ p.apply(Read.from(new TestUnboundedSource(TAG_SLIDING, 1, 6)))
+ .apply(ParDo.of(new ToKvFn()))
+ .setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()))
+ .apply(
+ Window.into(
+
SlidingWindows.of(Duration.standardSeconds(4)).every(Duration.standardSeconds(2))))
+ .apply(Count.perKey())
+ .apply(ParDo.of(new FormatCountFn()))
+ .apply(ParDo.of(new StreamingTestUtils.CollectDoFn<>(col)));
+
+ assertEquals(PipelineResult.State.DONE,
StreamingTestUtils.run(p).getState());
+ Set<String> res = StreamingTestUtils.collected(col);
+ assertTrue(res.contains((TestUnboundedSource.BASE_MILLIS - 2000) +
":k=2:ON_TIME"));
+ assertTrue(res.contains(TestUnboundedSource.BASE_MILLIS + ":k=4:ON_TIME"));
+ }
+
+ @Test
+ public void testCombineGloballyWithoutDefaults() throws Exception {
+ String col = StreamingTestUtils.newCollectorId(TAG_GLOBALLY);
+ Pipeline p = Pipeline.create(options(TAG_GLOBALLY));
+ p.apply(Read.from(new TestUnboundedSource(TAG_GLOBALLY, 1, 6)))
+ .apply(Window.into(FixedWindows.of(Duration.standardSeconds(2))))
+ .apply(Combine.globally(Count.<String>combineFn()).withoutDefaults())
+ .apply(ParDo.of(new FormatGlobalCountFn()))
+ .apply(ParDo.of(new StreamingTestUtils.CollectDoFn<>(col)));
+
+ assertEquals(PipelineResult.State.DONE,
StreamingTestUtils.run(p).getState());
+ Set<String> res = StreamingTestUtils.collected(col);
+ assertTrue(res.contains(TestUnboundedSource.BASE_MILLIS + "=2"));
+ assertTrue(res.contains((TestUnboundedSource.BASE_MILLIS + 2000) + "=2"));
+ }
+
+ @Test
+ public void testLateDataDropped() throws Exception {
+ String col = StreamingTestUtils.newCollectorId(TAG_LATE);
+ Pipeline p = Pipeline.create(options(TAG_LATE));
+ p.apply(Read.from(new TestUnboundedSource(TAG_LATE, 2, 8)))
+ .apply(ParDo.of(new ToKvFn()))
+ .setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()))
+ .apply(Window.into(FixedWindows.of(Duration.standardSeconds(4))))
+ .apply(Count.perKey())
+ .apply(ParDo.of(new FormatCountFn()))
+ .apply(ParDo.of(new StreamingTestUtils.CollectDoFn<>(col)));
+
+ assertEquals(PipelineResult.State.DONE,
StreamingTestUtils.run(p).getState());
+ Set<String> res = StreamingTestUtils.collected(col);
+ assertTrue(res.contains(TestUnboundedSource.BASE_MILLIS + ":k=1:ON_TIME"));
+ assertTrue(res.stream().noneMatch(s -> s.endsWith("LATE")));
+ }
+
+ @Test
+ public void testChainedStatefulParDoTimerOutputSurvives() throws Exception {
+ String col = StreamingTestUtils.newCollectorId(TAG_CHAINED);
+ Pipeline p = Pipeline.create(options(TAG_CHAINED));
+ p.apply(Read.from(new TestUnboundedSource(TAG_CHAINED, 1, 4)))
+ .apply(ParDo.of(new ToKvFn()))
+ .setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()))
+ .apply(ParDo.of(new TimerEmitterFn()))
+ .setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()))
+ .apply(Window.into(FixedWindows.of(Duration.millis(500))))
+ .apply(GroupByKey.create())
+ .apply(ParDo.of(new FormatGbkFn()))
+ .apply(ParDo.of(new StreamingTestUtils.CollectDoFn<>(col)));
+
+ assertEquals(PipelineResult.State.DONE,
StreamingTestUtils.run(p).getState());
+ Set<String> res = StreamingTestUtils.collected(col);
+ assertTrue(res.contains(TestUnboundedSource.BASE_MILLIS + ":chained=2"));
+ }
+
+ @Test
+ public void testExpiredKeyLeavesNoState() throws Exception {
+ String checkpoint = tempFolder.newFolder("cp").getAbsolutePath();
+ String col = StreamingTestUtils.newCollectorId(TAG_EXPIRY);
+ SparkStructuredStreamingPipelineOptions options =
+ StreamingTestUtils.streamingOptions(checkpoint);
+ options.setMaxRecordsPerBatch(1L);
+ Pipeline p = Pipeline.create(options);
+ p.apply(Read.from(new TestUnboundedSource(TAG_EXPIRY, 1, 8)))
+ .apply(ParDo.of(new ExpiryKeySelectorFn()))
+ .setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()))
+ .apply(Window.into(FixedWindows.of(Duration.standardSeconds(1))))
+ .apply(GroupByKey.create())
+ .apply(ParDo.of(new FormatGbkFn()))
+ .apply(ParDo.of(new StreamingTestUtils.CollectDoFn<>(col)));
+
+ assertEquals(PipelineResult.State.DONE,
StreamingTestUtils.run(p).getState());
+ Set<String> res = StreamingTestUtils.collected(col);
+ assertTrue(res.contains(TestUnboundedSource.BASE_MILLIS + ":k=1"));
+ for (String stateVarName : Arrays.asList("beamState", "beamTimers",
"beamWatermark")) {
+ assertEquals(stateVarName, 0L, StreamingTestUtils.stateRows(checkpoint,
stateVarName));
+ }
+ }
+
+ private static final class ToKvFn extends DoFn<String, KV<String, Integer>> {
+ @ProcessElement
+ public void process(@Element String el, OutputReceiver<KV<String,
Integer>> out) {
+ out.output(KV.of("k", TestUnboundedSource.indexOf(el)));
+ }
+ }
+
+ private static final class ExpiryKeySelectorFn extends DoFn<String,
KV<String, Integer>> {
+ @ProcessElement
+ public void process(@Element String el, OutputReceiver<KV<String,
Integer>> out) {
+ int i = TestUnboundedSource.indexOf(el);
+ if (i == 0) {
+ out.output(KV.of("k", i));
+ }
+ }
+ }
+
+ private static final class FormatGbkFn extends DoFn<KV<String,
Iterable<Integer>>, String> {
+ @ProcessElement
+ public void process(
+ @Element KV<String, Iterable<Integer>> el, IntervalWindow w,
OutputReceiver<String> out) {
+ out.output(w.start().getMillis() + ":" + el.getKey() + "=" +
Iterables.size(el.getValue()));
+ }
+ }
+
+ private static final class FormatCountFn extends DoFn<KV<String, Long>,
String> {
+ @ProcessElement
+ public void process(
+ @Element KV<String, Long> el, IntervalWindow w, PaneInfo pane,
OutputReceiver<String> out) {
+ out.output(
+ w.start().getMillis() + ":" + el.getKey() + "=" + el.getValue() +
":" + pane.getTiming());
+ }
+ }
+
+ private static final class FormatGlobalCountFn extends DoFn<Long, String> {
+ @ProcessElement
+ public void process(@Element Long count, IntervalWindow w,
OutputReceiver<String> out) {
+ out.output(w.start().getMillis() + "=" + count);
+ }
+ }
+
+ private static final class TimerEmitterFn extends DoFn<KV<String, Integer>,
KV<String, Integer>> {
+ @TimerId("t")
+ private final TimerSpec timer = TimerSpecs.timer(TimeDomain.EVENT_TIME);
+
+ @ProcessElement
+ public void process(
+ @Element KV<String, Integer> el,
+ @Timestamp Instant ts,
+ @TimerId("t") Timer t,
+ OutputReceiver<KV<String, Integer>> out) {
+ out.output(KV.of("chained", el.getValue()));
+ if (el.getValue() == 0) {
+ t.set(ts.plus(Duration.millis(250)));
+ }
+ }
+
+ @OnTimer("t")
+ public void onTimer(@Timestamp Instant ts, OutputReceiver<KV<String,
Integer>> out) {
+ out.outputWithTimestamp(KV.of("chained", 1), ts);
+ }
+ }
+}
