diff --git a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json index c0d5a77072b2..1b832d7f19cb 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 5a18b5da6fd8..6f2d14a69b8a 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 @@ 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.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, 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() { + @Override + protected void translate(GroupByKey transform, Context cxt) { + PCollection input = (PCollection) cxt.getInput(); + checkSupportedGroupByKey(input); super.translate(transform, cxt); } }; @@ -92,11 +105,17 @@ TransformTranslator getTransformTranslator(TransformT tra } 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 @@ private static void checkNoProcessingTimeTimers(DoFn doFn, DoFnSignature s } } - 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/GroupByKeyStreamingTranslator.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/GroupByKeyStreamingTranslator.java new file mode 100644 index 000000000000..b997f0d7498b --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/GroupByKeyStreamingTranslator.java @@ -0,0 +1,72 @@ +/* + * 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 org.apache.beam.runners.spark.coders.CoderHelpers; +import org.apache.beam.runners.spark.structuredstreaming.translation.TransformTranslator; +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.GroupByKey; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowingStrategy; +import org.apache.spark.api.java.function.MapFunction; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Encoders; +import org.apache.spark.sql.streaming.OutputMode; +import org.apache.spark.sql.streaming.TimeMode; + +/** Translates {@link GroupByKey} for Spark 4 Structured Streaming via transformWithState. */ +@Internal +public class GroupByKeyStreamingTranslator + extends TransformTranslator< + PCollection>, PCollection>>, GroupByKey> { + + public GroupByKeyStreamingTranslator() { + super(0.2f); + } + + @Override + protected void translate(GroupByKey transform, Context cxt) { + PCollection> input = cxt.getInput(); + WindowingStrategy windowing = input.getWindowingStrategy(); + KvCoder inputCoder = (KvCoder) input.getCoder(); + Coder keyCoder = inputCoder.getKeyCoder(); + PCollection>> output = cxt.getOutput(); + Coder>> outputCoder = output.getCoder(); + + GroupByKeyStatefulProcessor processor = + new GroupByKeyStatefulProcessor<>(inputCoder, windowing, cxt.getOptionsSupplier()); + + MapFunction>, byte[]> keyFn = + v -> CoderHelpers.toByteArray(v.getValue().getKey(), keyCoder); + Dataset>>> result = + cxt.getDataset(input) + .groupByKey(keyFn, Encoders.BINARY()) + .transformWithState( + processor, + TimeMode.EventTime(), + OutputMode.Append(), + cxt.windowedEncoder(outputCoder)); + + cxt.putDataset(output, result); + } +} 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 199570ae64a1..3f43a6b25e67 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.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 @@ protected void translate(ParDo.MultiOutput, OutputT> transform, Context cxt.getCurrentTransform(), input, SparkSideInputReader.empty(), false); MetricsAccumulator metrics = MetricsAccumulator.getInstance(cxt.getSparkSession()); - BeamStatefulProcessor processor = - new BeamStatefulProcessor<>( + ParDoStatefulProcessor processor = + new ParDoStatefulProcessor<>( doFn, runnerFactory, inputCoder, windowing, cxt.getOptionsSupplier(), metrics); MapFunction>, 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 7d794ef04506..568c41d23ec8 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.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.joda.time.Instant; import scala.collection.Iterator; -/** Runs a stateful {@link DoFn} within Spark 4 Structured Streaming {@code transformWithState}. */ -public class BeamStatefulProcessor +/** Base stateful processor bridging Beam state and timers over Spark transformWithState. */ +public abstract class BeamStatefulProcessor extends StatefulProcessor>, WindowedValue> { private static final String BEAM_STATE = "beamState"; private static final String BEAM_TIMERS = "beamTimers"; private static final String BEAM_WATERMARK = "beamWatermark"; - private final DoFn, OutputT> doFn; - private final DoFnRunnerFactory, OutputT> runnerFactory; - private final KvCoder inputCoder; - private final WindowingStrategy windowingStrategy; - private final Supplier optionsSupplier; - private final MetricsAccumulator metrics; - private final Coder windowCoder; + protected final KvCoder inputCoder; + protected final WindowingStrategy windowingStrategy; + protected final Coder windowCoder; private final TimerDataCoderV2 timerDataCoder; - private final StatefulTaskRunner, OutputT> taskRunner = new StatefulTaskRunner<>(); private transient @Nullable MapState beamState; private transient @Nullable ValueState beamTimers; private transient @Nullable ValueState beamWatermark; private transient @Nullable List> currentOutputs; - private transient boolean needsBundleStart; + + private transient @Nullable WindowedValueMultiReceiver outputReceiver; + + protected WindowedValueMultiReceiver outputReceiver() { + if (outputReceiver == null) { + outputReceiver = + new WindowedValueMultiReceiver() { + @Override + public void output(TupleTag tag, WindowedValue output) { + @SuppressWarnings("unchecked") + WindowedValue out = (WindowedValue) output; + checkStateNotNull(currentOutputs, "currentOutputs not initialized").add(out); + } + }; + } + return outputReceiver; + } public BeamStatefulProcessor( - DoFn, OutputT> doFn, - DoFnRunnerFactory, OutputT> runnerFactory, - KvCoder inputCoder, - WindowingStrategy windowingStrategy, - Supplier optionsSupplier, - MetricsAccumulator metrics) { - this.doFn = doFn; - this.runnerFactory = runnerFactory; + KvCoder inputCoder, WindowingStrategy windowingStrategy) { this.inputCoder = inputCoder; this.windowingStrategy = windowingStrategy; - this.optionsSupplier = optionsSupplier; - this.metrics = metrics; @SuppressWarnings("unchecked") Coder windowCoder = (Coder) windowingStrategy.getWindowFn().windowCoder(); @@ -151,16 +143,13 @@ public Iterator> handleInputRows( byte[] keyBytes, Iterator>> rows, TimerValues timerValues) { K key = CoderHelpers.fromByteArray(keyBytes, inputCoder.getKeyCoder()); List> outputs = new ArrayList<>(); + this.currentOutputs = outputs; SparkTimerInternals timerInternals = timerInternals(timerValues); + SparkStateInternals state = stateInternals(key); - DoFnRunner, 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 Iterator> handleExpiredTimer( byte[] keyBytes, TimerValues timerValues, ExpiredTimerInfo expiredTimerInfo) { K key = CoderHelpers.fromByteArray(keyBytes, inputCoder.getKeyCoder()); List> outputs = new ArrayList<>(); + this.currentOutputs = outputs; SparkTimerInternals timerInternals = timerInternals(timerValues); + SparkStateInternals state = stateInternals(key); timerInternals.advanceWatermark(); - DoFnRunner, 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 Iterator> handleExpiredTimer( return ScalaInterop.scalaIterator(outputs); } - @Override - public void close() { - taskRunner.teardownOnce(); + private SparkStateInternals 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, OutputT> baseRunner( - SparkStateInternals stateInternals, SparkTimerInternals timerInternals) { - return taskRunner.getOrCreate( - ctx -> { - ctx.set(stateInternals, timerInternals); - return runnerFactory.create( - optionsSupplier.get(), - metrics, - new WindowedValueMultiReceiver() { - @Override - public void output(TupleTag tag, WindowedValue output) { - @SuppressWarnings("unchecked") - WindowedValue out = (WindowedValue) output; - checkStateNotNull(currentOutputs, "currentOutputs not initialized").add(out); - } - }, - ctx); - }); + protected List 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, OutputT> keyRunner( - K key, SparkTimerInternals timerInternals, List> outputs) { - this.currentOutputs = outputs; - MapState state = checkStateNotNull(beamState); - SparkStateInternals stateInternals = - SparkStateInternals.forKey(key, new MapStateAdapter(state)); - DoFnRunnerWithTeardown, 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>> rows, + SparkTimerInternals timers, + SparkStateInternals state); - StatefulDoFnRunner.CleanupTimer> cleanupTimer = - new StatefulDoFnRunner.TimeInternalsCleanupTimer>( - timerInternals, windowingStrategy) { - @Override - public void setForWindow(KV 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 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 state); private void restoreTimers(SparkTimerInternals timerInternals) { ValueState timersState = checkStateNotNull(beamTimers); @@ -278,28 +223,6 @@ private void persistTimers(SparkTimerInternals timerInternals) { } } - private void fireDueTimers( - K key, - Instant watermark, - SparkTimerInternals timerInternals, - DoFnRunner, 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 000000000000..cbe592af0eae --- /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 + extends BeamStatefulProcessor>> { + + private final Supplier optionsSupplier; + private final TupleTag>> mainTag = new TupleTag<>("output"); + + private transient @Nullable DoFnRunner, KV>> runner; + private transient @Nullable MutableStepContext stepContext; + + public GroupByKeyStatefulProcessor( + KvCoder inputCoder, + WindowingStrategy windowingStrategy, + Supplier optionsSupplier) { + super(inputCoder, windowingStrategy); + this.optionsSupplier = optionsSupplier; + } + + @Override + protected void processElements( + K key, + Iterator>> rows, + SparkTimerInternals timerInternals, + SparkStateInternals stateInternals) { + DoFnRunner, KV>> r = runner(); + checkStateNotNull(stepContext).set(stateInternals, timerInternals); + + List> elementValues = new ArrayList<>(); + while (rows.hasNext()) { + WindowedValue> 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 stateInternals) { + List dueTimers = dueEventTimeTimers(timerInternals); + if (dueTimers.isEmpty()) { + return; + } + for (TimerData timer : dueTimers) { + timerInternals.deleteTimer(timer); + } + DoFnRunner, KV>> r = runner(); + checkStateNotNull(stepContext).set(stateInternals, timerInternals); + r.processElement( + WindowedValues.valueInGlobalWindow(KeyedWorkItems.timersWorkItem(key, dueTimers))); + } + + private DoFnRunner, KV>> runner() { + if (runner == null) { + MutableStepContext ctx = new MutableStepContext(); + this.stepContext = ctx; + WindowedValueMultiReceiver receiver = outputReceiver(); + DoFn, KV>> gabwDoFn = + GroupAlsoByWindowViaWindowSetNewDoFn.create( + windowingStrategy, + k -> ctx.stateInternals(), + k -> ctx.timerInternals(), + NullSideInputReader.empty(), + SystemReduceFn.buffering(inputCoder.getValueCoder()), + receiver, + mainTag); + + Coder> itemCoder = + KeyedWorkItemCoder.of(inputCoder.getKeyCoder(), inputCoder.getValueCoder(), windowCoder); + + DoFnRunner, KV>> 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 000000000000..f25c523eaac6 --- /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 extends BeamStatefulProcessor { + + private final DoFn, OutputT> doFn; + private final DoFnRunnerFactory, OutputT> runnerFactory; + private final Supplier optionsSupplier; + private final MetricsAccumulator metrics; + private final StatefulTaskRunner, OutputT> taskRunner = new StatefulTaskRunner<>(); + + private transient boolean needsBundleStart; + + public ParDoStatefulProcessor( + DoFn, OutputT> doFn, + DoFnRunnerFactory, OutputT> runnerFactory, + KvCoder inputCoder, + WindowingStrategy windowingStrategy, + Supplier 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>> rows, + SparkTimerInternals timerInternals, + SparkStateInternals stateInternals) { + DoFnRunner, 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 stateInternals) { + if (dueEventTimeTimers(timerInternals).isEmpty()) { + return; + } + DoFnRunner, 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, OutputT> runner) { + while (true) { + List due = dueEventTimeTimers(timerInternals); + if (due.isEmpty()) { + break; + } + TimerData next = due.get(0); + timerInternals.deleteTimer(next); + StatefulTaskRunner.fireTimer(runner, key, next); + } + } + + private DoFnRunnerWithTeardown, OutputT> baseRunner( + SparkStateInternals stateInternals, SparkTimerInternals timerInternals) { + return taskRunner.getOrCreate( + ctx -> { + ctx.set(stateInternals, timerInternals); + return runnerFactory.create(optionsSupplier.get(), metrics, outputReceiver(), ctx); + }); + } + + private DoFnRunner, OutputT> keyRunner( + SparkStateInternals stateInternals, SparkTimerInternals timerInternals) { + DoFnRunnerWithTeardown, OutputT> base = baseRunner(stateInternals, timerInternals); + taskRunner.stepContext().set(stateInternals, timerInternals); + + // Closed in processElements and processTimers. + if (needsBundleStart) { + needsBundleStart = false; + base.startBundle(); + } + + StatefulDoFnRunner.CleanupTimer> cleanupTimer = + new StatefulDoFnRunner.TimeInternalsCleanupTimer>( + timerInternals, windowingStrategy) { + @Override + public void setForWindow(KV 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 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 39992b76fe3e..a5afa7c1a4a6 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.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 @@ private static void assertUnsupported(Pipeline pipeline, String expected) { 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 void rejectsUnsupportedStatefulFeatures() throws Exception { .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.>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.>into(FixedWindows.of(Duration.millis(1))) + .accumulatingFiredPanes()) + .apply(GroupByKey.create()) + .getPipeline(), + "Accumulating fired panes" + }, + { + kv("cgdef") + .apply( + Window.>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 18143ce4d924..5679950e1893 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.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 void testExpiredKeyLeavesNoState() throws Exception { Set 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> { @ProcessElement public void process(@Element String element, OutputReceiver> 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 3d7841154150..5616b5b4b7c8 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.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 static SparkStructuredStreamingPipelineOptions streamingOptions(String ch 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 000000000000..2b0305fa66d2 --- /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 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 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.combineFn()).withoutDefaults()) + .apply(ParDo.of(new FormatGlobalCountFn())) + .apply(ParDo.of(new StreamingTestUtils.CollectDoFn<>(col))); + + assertEquals(PipelineResult.State.DONE, StreamingTestUtils.run(p).getState()); + Set 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 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 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 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> { + @ProcessElement + public void process(@Element String el, OutputReceiver> out) { + out.output(KV.of("k", TestUnboundedSource.indexOf(el))); + } + } + + private static final class ExpiryKeySelectorFn extends DoFn> { + @ProcessElement + public void process(@Element String el, OutputReceiver> out) { + int i = TestUnboundedSource.indexOf(el); + if (i == 0) { + out.output(KV.of("k", i)); + } + } + } + + private static final class FormatGbkFn extends DoFn>, String> { + @ProcessElement + public void process( + @Element KV> el, IntervalWindow w, OutputReceiver out) { + out.output(w.start().getMillis() + ":" + el.getKey() + "=" + Iterables.size(el.getValue())); + } + } + + private static final class FormatCountFn extends DoFn, String> { + @ProcessElement + public void process( + @Element KV el, IntervalWindow w, PaneInfo pane, OutputReceiver out) { + out.output( + w.start().getMillis() + ":" + el.getKey() + "=" + el.getValue() + ":" + pane.getTiming()); + } + } + + private static final class FormatGlobalCountFn extends DoFn { + @ProcessElement + public void process(@Element Long count, IntervalWindow w, OutputReceiver out) { + out.output(w.start().getMillis() + "=" + count); + } + } + + private static final class TimerEmitterFn extends DoFn, KV> { + @TimerId("t") + private final TimerSpec timer = TimerSpecs.timer(TimeDomain.EVENT_TIME); + + @ProcessElement + public void process( + @Element KV el, + @Timestamp Instant ts, + @TimerId("t") Timer t, + OutputReceiver> 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> out) { + out.outputWithTimestamp(KV.of("chained", 1), ts); + } + } +}