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 f8221572fb2 [Spark][#36841] Translate stateless streaming pipelines on
the Spark 4 runner (#40090)
f8221572fb2 is described below
commit f8221572fb20b18eda0530bca9f117eb36c68e1d
Author: Tobias Kaymak <[email protected]>
AuthorDate: Fri Sep 11 20:42:24 2026 +0200
[Spark][#36841] Translate stateless streaming pipelines on the Spark 4
runner (#40090)
* [Spark 4] Translate stateless streaming pipelines
Makes the DataSourceV2 unbounded source from #39971 reachable. The Spark 4
module overrides PipelineTranslatorFactory and dispatches streaming
pipelines
to PipelineTranslatorStreaming, which translates unbounded reads and reuses
the batch translators for stateless single output ParDo, Window.Assign,
Flatten and Reshuffle. GroupByKey, Combine.perKey, stateful ParDo, ParDo
with
side inputs or additional outputs, Impulse and bounded reads fail at
translation, the batch translators for them persist or collect the Dataset,
which Spark rejects on a streaming plan.
---
.../translation/PipelineTranslatorFactory.java | 35 +++
.../translation/PipelineTranslatorStreaming.java | 113 +++++++++
.../translation/StreamingEvaluationContext.java | 261 ++++++++++++++++++++
.../streaming/ReadUnboundedTranslator.java | 97 ++++++++
.../io/streaming/BeamMicroBatchSourceTest.java | 269 ++-------------------
.../io/streaming/TestUnboundedSource.java | 252 +++++++++++++++++++
.../PipelineTranslatorStreamingTest.java | 142 +++++++++++
.../streaming/StatelessParDoStreamingTest.java | 110 +++++++++
.../streaming/StreamingCheckpointRestartTest.java | 118 +++++++++
.../streaming/StreamingPipelineLifecycleTest.java | 171 +++++++++++++
.../translation/streaming/StreamingTestUtils.java | 156 ++++++++++++
.../translation/batch/ParDoTranslatorBatch.java | 2 +-
.../translation/batch/PipelineTranslatorBatch.java | 91 +------
...torBatch.java => PipelineTranslatorCommon.java} | 8 +-
.../batch/StatefulParDoTranslatorBatch.java | 2 +-
15 files changed, 1486 insertions(+), 341 deletions(-)
diff --git
a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorFactory.java
b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorFactory.java
new file mode 100644
index 00000000000..9f1c15a4777
--- /dev/null
+++
b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorFactory.java
@@ -0,0 +1,35 @@
+/*
+ * 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;
+
+import
org.apache.beam.runners.spark.structuredstreaming.translation.batch.PipelineTranslatorBatch;
+import org.apache.beam.sdk.annotations.Internal;
+
+/**
+ * This class shadows the shared base file of the same name. The Spark 4
module compiles the
+ * override tree with later wins, so this copy replaces the base one that
throws for streaming.
+ */
+@Internal
+public final class PipelineTranslatorFactory {
+ private PipelineTranslatorFactory() {}
+
+ /** Creates a {@link PipelineTranslator} for the given execution mode. */
+ public static PipelineTranslator create(boolean streaming) {
+ return streaming ? new PipelineTranslatorStreaming() : new
PipelineTranslatorBatch();
+ }
+}
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
new file mode 100644
index 00000000000..3529daaff31
--- /dev/null
+++
b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java
@@ -0,0 +1,113 @@
+/*
+ * 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;
+
+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.ReadUnboundedTranslator;
+import org.apache.beam.sdk.annotations.Internal;
+import org.apache.beam.sdk.transforms.Combine;
+import org.apache.beam.sdk.transforms.GroupByKey;
+import org.apache.beam.sdk.transforms.Impulse;
+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.util.construction.SplittableParDo;
+import org.apache.beam.sdk.values.PInput;
+import org.apache.beam.sdk.values.POutput;
+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.
+ */
+@Internal
+public class PipelineTranslatorStreaming extends PipelineTranslatorCommon {
+
+ @SuppressWarnings("rawtypes")
+ private static final TransformTranslator READ_UNBOUNDED = new
ReadUnboundedTranslator<>();
+
+ private static final String NOT_SUPPORTED =
+ " is not supported by the Spark 4 streaming runner yet, see"
+ + " https://github.com/apache/beam/issues/36841";;
+
+ /** Returns a {@link TransformTranslator} for the given {@link PTransform}
if known. */
+ @Override
+ @Nullable
+ protected <InT extends PInput, OutT extends POutput, TransformT extends
PTransform<InT, OutT>>
+ TransformTranslator<InT, OutT, TransformT>
getTransformTranslator(TransformT transform) {
+
+ if (transform instanceof SplittableParDo.PrimitiveUnboundedRead) {
+ return READ_UNBOUNDED;
+ }
+
+ if (transform instanceof SplittableParDo.PrimitiveBoundedRead) {
+ throw new UnsupportedOperationException(
+ "Bounded Read (Read.from(BoundedSource), Create with two or more
elements)"
+ + NOT_SUPPORTED);
+ }
+
+ if (transform instanceof Impulse) {
+ throw new UnsupportedOperationException(
+ "Impulse (Create with fewer than two elements, PAssert)" +
NOT_SUPPORTED);
+ }
+
+ if (transform instanceof GroupByKey) {
+ throw new UnsupportedOperationException("GroupByKey" + NOT_SUPPORTED);
+ }
+
+ if (transform instanceof Combine.PerKey) {
+ throw new UnsupportedOperationException("Combine.perKey" +
NOT_SUPPORTED);
+ }
+
+ if (transform instanceof ParDo.MultiOutput) {
+ ParDo.MultiOutput<?, ?> parDo = (ParDo.MultiOutput<?, ?>) transform;
+ DoFnSignature signature = DoFnSignatures.signatureForDoFn(parDo.getFn());
+ if (signature.usesState() || signature.usesTimers()) {
+ throw new UnsupportedOperationException(
+ "Stateful ParDo (" + signature.fnClass().getName() + ")" +
NOT_SUPPORTED);
+ }
+ if (!parDo.getSideInputs().isEmpty()) {
+ throw new UnsupportedOperationException(
+ "ParDo with side inputs (" + signature.fnClass().getName() + ")" +
NOT_SUPPORTED);
+ }
+ if (!parDo.getAdditionalOutputTags().getAll().isEmpty()) {
+ throw new UnsupportedOperationException(
+ "ParDo with additional outputs ("
+ + signature.fnClass().getName()
+ + ")"
+ + NOT_SUPPORTED);
+ }
+ }
+
+ return super.getTransformTranslator(transform);
+ }
+
+ @Override
+ protected EvaluationContext createEvaluationContext(
+ Collection<? extends EvaluationContext.NamedDataset<?>> leaves,
+ SparkSession session,
+ SparkCommonPipelineOptions options) {
+ return new StreamingEvaluationContext(leaves, session, options);
+ }
+}
diff --git
a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/StreamingEvaluationContext.java
b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/StreamingEvaluationContext.java
new file mode 100644
index 00000000000..ff7839bddb2
--- /dev/null
+++
b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/StreamingEvaluationContext.java
@@ -0,0 +1,261 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.spark.structuredstreaming.translation;
+
+import static
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicInteger;
+import org.apache.beam.runners.spark.SparkCommonPipelineOptions;
+import
org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions;
+import org.apache.beam.sdk.annotations.Internal;
+import org.apache.hadoop.fs.Path;
+import org.apache.spark.sql.Dataset;
+import org.apache.spark.sql.SparkSession;
+import org.apache.spark.sql.streaming.StreamingQuery;
+import org.apache.spark.sql.streaming.StreamingQueryException;
+import org.apache.spark.sql.streaming.StreamingQueryListener;
+import org.apache.spark.sql.streaming.Trigger;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Starts one Spark Structured Streaming query per leaf dataset and blocks
until all of them reach a
+ * terminal state. Queries end through {@link #stop()} or the idle stop
listener.
+ *
+ * <p>Leaf {@code i} checkpoints under {@code <checkpointDir>/i}, in pipeline
graph order. A changed
+ * pipeline needs a new checkpoint directory, as with any Spark streaming
query.
+ */
+@Internal
+public class StreamingEvaluationContext extends EvaluationContext {
+ private static final Logger LOG =
LoggerFactory.getLogger(StreamingEvaluationContext.class);
+
+ private static final long AWAIT_POLL_TIMEOUT_MILLIS = 100;
+
+ private final SparkStructuredStreamingPipelineOptions options;
+
+ // Guards queries and stopped.
+ private final Object lock = new Object();
+ private final List<StreamingQuery> queries = new ArrayList<>();
+ private boolean stopped = false;
+
+ StreamingEvaluationContext(
+ Collection<? extends NamedDataset<?>> leaves,
+ SparkSession session,
+ SparkCommonPipelineOptions options) {
+ super(leaves, session);
+ this.options = options.as(SparkStructuredStreamingPipelineOptions.class);
+ }
+
+ /** Starts one streaming query per leaf dataset and blocks until all queries
terminate. */
+ @Override
+ public void evaluate() {
+ String checkpointBaseDir = options.getCheckpointDir();
+ checkArgument(
+ checkpointBaseDir != null && !checkpointBaseDir.isEmpty(),
+ "checkpointDir must be set for a streaming pipeline");
+ int idleStopThreshold = options.getStreamingStopAfterIdleBatches();
+
+ StreamingQueryListener idleStopListener = null;
+ if (idleStopThreshold >= 0) {
+ idleStopListener = new IdleStopListener(idleStopThreshold);
+ getSparkSession().streams().addListener(idleStopListener);
+ }
+
+ try {
+ int leafIndex = 0;
+ for (NamedDataset<?> ds : leaves()) {
+ Dataset<?> dataset = ds.dataset();
+ if (dataset == null) {
+ continue;
+ }
+ synchronized (lock) {
+ if (stopped) {
+ break;
+ }
+ }
+ if (!dataset.isStreaming()) {
+ EvaluationContext.evaluate(ds.name(), dataset);
+ continue;
+ }
+
+ StreamingQuery query = startQuery(dataset, checkpointBaseDir,
leafIndex++, options);
+ boolean alreadyStopped;
+ synchronized (lock) {
+ queries.add(query);
+ alreadyStopped = stopped;
+ }
+ if (alreadyStopped) {
+ stopQuery(query);
+ }
+ }
+
+ List<StreamingQuery> toAwait;
+ synchronized (lock) {
+ toAwait = new ArrayList<>(queries);
+ }
+ awaitTermination(toAwait);
+ } catch (RuntimeException e) {
+ stop();
+ throw e;
+ } finally {
+ if (idleStopListener != null) {
+ getSparkSession().streams().removeListener(idleStopListener);
+ }
+ }
+ }
+
+ /**
+ * Stops all queries started by {@link #evaluate()}. This method is
idempotent and thread safe.
+ */
+ @Override
+ public void stop() {
+ List<StreamingQuery> toStop;
+ synchronized (lock) {
+ if (stopped) {
+ return;
+ }
+ stopped = true;
+ toStop = new ArrayList<>(queries);
+ }
+ for (StreamingQuery query : toStop) {
+ stopQuery(query);
+ }
+ }
+
+ private StreamingQuery startQuery(
+ Dataset<?> dataset,
+ String checkpointBaseDir,
+ int leafIndex,
+ SparkStructuredStreamingPipelineOptions options) {
+ try {
+ return dataset
+ .writeStream()
+ .format("noop")
+ .outputMode("append")
+ .option(
+ "checkpointLocation",
+ new Path(checkpointBaseDir,
Integer.toString(leafIndex)).toString())
+ .trigger(Trigger.ProcessingTime(options.getMaxBatchDurationMillis()))
+ .start();
+ } catch (TimeoutException e) {
+ throw new RuntimeException(
+ "Failed to start streaming query for leaf dataset index " +
leafIndex, e);
+ }
+ }
+
+ /** Blocks until every query in toAwait has terminated. Sibling queries stop
on failure. */
+ private void awaitTermination(List<StreamingQuery> toAwait) {
+ List<StreamingQuery> active = new ArrayList<>(toAwait);
+ while (!active.isEmpty()) {
+ Iterator<StreamingQuery> iterator = active.iterator();
+ while (iterator.hasNext()) {
+ StreamingQuery query = iterator.next();
+ try {
+ if (query.awaitTermination(AWAIT_POLL_TIMEOUT_MILLIS)) {
+ iterator.remove();
+ }
+ } catch (StreamingQueryException e) {
+ LOG.error("Streaming query {} terminated with an exception.",
query.id(), e);
+ stop();
+ throw new RuntimeException(e);
+ }
+ }
+ }
+ }
+
+ /** Stops a single query if active. */
+ private void stopQuery(StreamingQuery query) {
+ try {
+ if (query.isActive()) {
+ query.stop();
+ }
+ } catch (TimeoutException | RuntimeException e) {
+ LOG.warn("Failed to stop streaming query {}.", query.id(), e);
+ }
+ }
+
+ private void stopQueryById(UUID id) {
+ StreamingQuery match = null;
+ synchronized (lock) {
+ for (StreamingQuery query : queries) {
+ if (query.id().equals(id)) {
+ match = query;
+ break;
+ }
+ }
+ }
+ if (match != null) {
+ stopQuery(match);
+ }
+ }
+
+ /**
+ * Stops a query after {@code threshold} consecutive triggers without input
rows. A trigger
+ * without data is reported as a progress event with zero rows when the
source offset moved and as
+ * an idle event otherwise, both count.
+ */
+ private final class IdleStopListener extends StreamingQueryListener {
+ private final int threshold;
+ private final Map<UUID, AtomicInteger> idleCounts = new
ConcurrentHashMap<>();
+
+ IdleStopListener(int threshold) {
+ this.threshold = threshold;
+ }
+
+ @Override
+ public void onQueryStarted(QueryStartedEvent event) {}
+
+ @Override
+ public void onQueryProgress(QueryProgressEvent event) {
+ UUID id = event.progress().id();
+ if (event.progress().numInputRows() == 0) {
+ countIdle(id);
+ } else {
+ idleCounts.remove(id);
+ }
+ }
+
+ @Override
+ public void onQueryIdle(QueryIdleEvent event) {
+ countIdle(event.id());
+ }
+
+ private void countIdle(UUID id) {
+ int count = idleCounts.computeIfAbsent(id, unused -> new
AtomicInteger()).incrementAndGet();
+ if (count >= threshold) {
+ idleCounts.remove(id);
+ Thread stopThread = new Thread(() -> stopQueryById(id),
"beam-idle-stop-" + id);
+ stopThread.setDaemon(true);
+ stopThread.start();
+ }
+ }
+
+ @Override
+ public void onQueryTerminated(QueryTerminatedEvent event) {
+ idleCounts.remove(event.id());
+ }
+ }
+}
diff --git
a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ReadUnboundedTranslator.java
b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ReadUnboundedTranslator.java
new file mode 100644
index 00000000000..fb1691108b7
--- /dev/null
+++
b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ReadUnboundedTranslator.java
@@ -0,0 +1,97 @@
+/*
+ * 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.apache.spark.sql.functions.col;
+
+import
org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions;
+import
org.apache.beam.runners.spark.structuredstreaming.io.streaming.UnboundedSourceDataset;
+import
org.apache.beam.runners.spark.structuredstreaming.translation.TransformTranslator;
+import
org.apache.beam.runners.spark.structuredstreaming.translation.helpers.CoderHelpers;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.io.UnboundedSource;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.util.construction.SplittableParDo;
+import org.apache.beam.sdk.values.PBegin;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.WindowedValue;
+import org.apache.beam.sdk.values.WindowedValues;
+import org.apache.spark.api.java.function.MapFunction;
+import org.apache.spark.sql.Dataset;
+import org.apache.spark.sql.Encoder;
+import org.apache.spark.sql.Encoders;
+import org.apache.spark.sql.Row;
+
+/**
+ * Translator for {@link SplittableParDo.PrimitiveUnboundedRead}.
+ *
+ * <p>Elements arrive in the global window with the record timestamp.
Downstream windowing requires
+ * an explicit {@code Window.Assign}.
+ */
+public class ReadUnboundedTranslator<T>
+ extends TransformTranslator<PBegin, PCollection<T>,
SplittableParDo.PrimitiveUnboundedRead<T>> {
+
+ public ReadUnboundedTranslator() {
+ super(0.05f);
+ }
+
+ @Override
+ protected void translate(SplittableParDo.PrimitiveUnboundedRead<T>
transform, Context cxt) {
+ PCollection<T> output = cxt.getOutput();
+ UnboundedSource<T, ?> source = transform.getSource();
+ Coder<T> elementCoder = output.getCoder();
+
+ WindowedValues.FullWindowedValueCoder<T> payloadCoder =
+ WindowedValues.getFullCoder(elementCoder, GlobalWindow.Coder.INSTANCE);
+
+ SparkStructuredStreamingPipelineOptions options =
+ cxt.getOptions().as(SparkStructuredStreamingPipelineOptions.class);
+
+ Dataset<Row> rows =
+ UnboundedSourceDataset.of(
+ cxt.getSparkSession(),
+ source,
+ payloadCoder,
+ options,
+ cxt.getCurrentTransform().getFullName());
+
+ Encoder<WindowedValue<T>> encoder =
+ cxt.windowedEncoder(elementCoder, GlobalWindow.Coder.INSTANCE);
+
+ Dataset<WindowedValue<T>> dataset =
+ rows.select(col(UnboundedSourceDataset.COL_PAYLOAD))
+ .as(Encoders.BINARY())
+ .map(new DecodePayload<>(payloadCoder), encoder);
+
+ cxt.putDataset(output, dataset);
+ }
+
+ /** Decodes the binary payload column back into a Beam {@code
WindowedValue}. */
+ private static final class DecodePayload<T> implements MapFunction<byte[],
WindowedValue<T>> {
+ private final Coder<WindowedValue<T>> coder;
+
+ DecodePayload(Coder<WindowedValue<T>> coder) {
+ this.coder = coder;
+ }
+
+ @Override
+ public WindowedValue<T> call(byte[] payload) {
+ return CoderHelpers.fromByteArray(payload, coder);
+ }
+ }
+}
diff --git
a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java
b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java
index 4616c7563a2..3b2e36bf5ef 100644
---
a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java
+++
b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java
@@ -28,8 +28,6 @@ import static org.junit.Assert.assertTrue;
import java.io.File;
import java.io.IOException;
-import java.io.InputStream;
-import java.io.OutputStream;
import java.io.Serializable;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
@@ -40,11 +38,9 @@ import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
-import java.util.NoSuchElementException;
import java.util.Set;
import java.util.TreeSet;
import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.BooleanSupplier;
import javax.annotation.Nullable;
@@ -58,15 +54,11 @@ import
org.apache.beam.runners.spark.structuredstreaming.io.streaming.UnboundedS
import
org.apache.beam.runners.spark.structuredstreaming.io.streaming.UnboundedSourceDataset.BeamPartitionReader;
import
org.apache.beam.runners.spark.structuredstreaming.io.streaming.UnboundedSourceDataset.BeamTable;
import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.sdk.coders.CustomCoder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
-import org.apache.beam.sdk.coders.VarIntCoder;
import org.apache.beam.sdk.coders.VarLongCoder;
import org.apache.beam.sdk.io.CountingSource;
import org.apache.beam.sdk.io.UnboundedSource;
-import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
-import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
import org.apache.beam.sdk.util.CoderUtils;
import org.apache.beam.sdk.values.WindowedValue;
@@ -88,7 +80,6 @@ import org.apache.spark.sql.streaming.StreamingQueryProgress;
import org.apache.spark.sql.streaming.Trigger;
import org.apache.spark.sql.util.CaseInsensitiveStringMap;
import org.apache.spark.util.SerializableConfiguration;
-import org.joda.time.Instant;
import org.junit.After;
import org.junit.Before;
import org.junit.BeforeClass;
@@ -124,11 +115,6 @@ public class BeamMicroBatchSourceTest implements
Serializable {
private static final Coder<WindowedValue<String>> CODER =
WindowedValues.getFullCoder(StringUtf8Coder.of(),
GlobalWindow.Coder.INSTANCE);
- /** 2023-11-14T22:13:20Z, a plain modern timestamp with no rebase or DST
subtleties. */
- private static final long BASE_MILLIS = 1_700_000_000_000L;
-
- private static final long INTERVAL_MILLIS = 1_000L;
-
private static final long POLL_TIMEOUT_MILLIS = 120_000L;
private static Broadcast<SerializablePipelineOptions> optionsBroadcast;
@@ -164,7 +150,7 @@ public class BeamMicroBatchSourceTest implements
Serializable {
public void tearDown() {
BeamReaderCache.invalidateAll();
BATCHES.clear();
- TestSource.forget(tag);
+ TestUnboundedSource.forget(tag);
}
/** The {@code EventTimeWatermark} node survives typed maps in the logical
and analyzed plan. */
@@ -230,12 +216,14 @@ public class BeamMicroBatchSourceTest implements
Serializable {
assertEquals(
Collections.singletonList(GlobalWindow.INSTANCE), new
ArrayList<>(value.getWindows()));
assertEquals(
- BASE_MILLIS + TestSource.indexOf(value.getValue()) *
INTERVAL_MILLIS,
+ TestUnboundedSource.BASE_MILLIS
+ + TestUnboundedSource.indexOf(value.getValue())
+ * TestUnboundedSource.INTERVAL_MILLIS,
value.getTimestamp().getMillis());
}
}
assertEquals(count, values.size());
- assertEquals(TestSource.elements(tag, 1, count), new HashSet<>(values));
+ assertEquals(TestUnboundedSource.elements(tag, 1, count), new
HashSet<>(values));
}
/** The default record limit is unlimited, an available source drains in one
micro-batch. */
@@ -315,7 +303,8 @@ public class BeamMicroBatchSourceTest implements
Serializable {
List<Integer> sizes = nonEmptySizes(batches(tag));
assertFalse("no rows arrived", sizes.isEmpty());
assertTrue("batch exceeds the shared limit: " + sizes,
Collections.max(sizes) <= 10);
- assertEquals(TestSource.elements(tag, shards, count), new
HashSet<>(values(batches(tag))));
+ assertEquals(
+ TestUnboundedSource.elements(tag, shards, count), new
HashSet<>(values(batches(tag))));
}
/**
@@ -353,7 +342,7 @@ public class BeamMicroBatchSourceTest implements
Serializable {
BeamReaderCache.invalidateAll();
List<String> firstValues = values(batches(first));
- Set<String> all = TestSource.elements(tag, shards, count);
+ Set<String> all = TestUnboundedSource.elements(tag, shards, count);
query = start(rows(shards, count, limited(limit, 1_000L)), second,
checkpointDir);
try {
await(
@@ -377,8 +366,8 @@ public class BeamMicroBatchSourceTest implements
Serializable {
for (int shard = 0; shard < shards; shard++) {
int min = Integer.MAX_VALUE;
for (String value : secondValues) {
- if (TestSource.shardOf(value) == shard) {
- min = Math.min(min, TestSource.indexOf(value));
+ if (TestUnboundedSource.shardOf(value) == shard) {
+ min = Math.min(min, TestUnboundedSource.indexOf(value));
}
}
assertTrue("run 2 delivered nothing for shard " + shard, min <
Integer.MAX_VALUE);
@@ -395,7 +384,7 @@ public class BeamMicroBatchSourceTest implements
Serializable {
int finalizations = 0;
for (int shard = 0; shard < 2; shard++) {
int committed = committedPosition(checkpointDir, shard);
- List<Integer> finalized = TestSource.finalized(tag, shard);
+ List<Integer> finalized = TestUnboundedSource.finalized(tag, shard);
assertTrue(
"shard " + shard + " finalized " + finalized + " beyond committed "
+ committed,
finalized.isEmpty() || Collections.max(finalized) <= committed);
@@ -460,7 +449,7 @@ public class BeamMicroBatchSourceTest implements
Serializable {
for (int shard = 0; shard < 2; shard++) {
int committed = committedPosition(checkpointDir, shard);
- List<Integer> finalized = TestSource.finalized(tag, shard);
+ List<Integer> finalized = TestUnboundedSource.finalized(tag, shard);
assertTrue(
"shard " + shard + " finalized " + finalized + ", committed " +
committed,
finalized.contains(committed) && Collections.max(finalized) ==
committed);
@@ -473,8 +462,8 @@ public class BeamMicroBatchSourceTest implements
Serializable {
String location = sourceDir(temp.newFolder("protocol")).getAbsolutePath();
assertEquals(shardZero(0, 1, 2), readBatch(partition(location, 0, 1)));
assertEquals(shardZero(0, 1, 2), readBatch(partition(location, 0, 1)));
- assertEquals(Collections.emptyList(), TestSource.finalized(tag, 0));
- assertEquals(2, TestSource.created(tag));
+ assertEquals(Collections.emptyList(), TestUnboundedSource.finalized(tag,
0));
+ assertEquals(2, TestUnboundedSource.created(tag));
}
/** A start epoch above zero without a durable mark is an invariant
violation. */
@@ -483,7 +472,7 @@ public class BeamMicroBatchSourceTest implements
Serializable {
String location = sourceDir(temp.newFolder("protocol")).getAbsolutePath();
assertThrows(
IllegalStateException.class, () -> new
BeamPartitionReader<>(partition(location, 5, 6)));
- assertEquals(0, TestSource.created(tag));
+ assertEquals(0, TestUnboundedSource.created(tag));
}
/** A failed mark write fails the batch after its rows, the retry recreates
the reader. */
@@ -494,8 +483,8 @@ public class BeamMicroBatchSourceTest implements
Serializable {
String file = location.getAbsolutePath();
assertEquals(shardZero(0, 1, 2), drainUntilFailure(partition(file, 0, 1),
IOException.class));
assertEquals(shardZero(0, 1, 2), drainUntilFailure(partition(file, 0, 1),
IOException.class));
- assertEquals(Collections.emptyList(), TestSource.finalized(tag, 0));
- assertEquals(2, TestSource.created(tag));
+ assertEquals(Collections.emptyList(), TestUnboundedSource.finalized(tag,
0));
+ assertEquals(2, TestUnboundedSource.created(tag));
}
//
---------------------------------------------------------------------------------------------
@@ -521,10 +510,10 @@ public class BeamMicroBatchSourceTest implements
Serializable {
int shards, int count, SparkStructuredStreamingPipelineOptions options) {
return UnboundedSourceDataset.of(
SESSION.getSession(),
- new TestSource(tag, shards, count),
+ new TestUnboundedSource(tag, shards, count),
CODER,
options,
- "Read(TestSource)");
+ "Read(TestUnboundedSource)");
}
/** Builds the driver side stream through the table, with the session's
broadcasts. */
@@ -624,7 +613,7 @@ public class BeamMicroBatchSourceTest implements
Serializable {
private static Set<Integer> shardsOf(List<String> values) {
Set<Integer> shards = new HashSet<>();
for (String value : values) {
- shards.add(TestSource.shardOf(value));
+ shards.add(TestUnboundedSource.shardOf(value));
}
return shards;
}
@@ -759,7 +748,7 @@ public class BeamMicroBatchSourceTest implements
Serializable {
new BeamSourceCheckpoint(sourceDir(checkpointDir).getAbsolutePath(),
new Configuration());
byte[] coded = checkpoint.readMark(shard, epoch);
assertNotNull("no mark at committed epoch " + epoch + " for shard " +
shard, coded);
- return CoderUtils.decodeFromByteArray(TestSource.MARK_CODER, coded).next;
+ return CoderUtils.decodeFromByteArray(TestUnboundedSource.MARK_CODER,
coded).next;
}
//
---------------------------------------------------------------------------------------------
@@ -768,7 +757,8 @@ public class BeamMicroBatchSourceTest implements
Serializable {
/** Split 0 of a single shard source of 100 elements from epoch {@code
start} to {@code end}. */
private BeamInputPartition<String> partition(String location, long start,
long end) {
- TestSource split = new TestSource(tag, 1, 100).split(1,
PipelineOptionsFactory.create()).get(0);
+ TestUnboundedSource split =
+ new TestUnboundedSource(tag, 1, 100).split(1,
PipelineOptionsFactory.create()).get(0);
return new BeamInputPartition<>(
split,
CODER,
@@ -810,219 +800,8 @@ public class BeamMicroBatchSourceTest implements
Serializable {
private List<String> shardZero(int... indexes) {
List<String> elements = new ArrayList<>();
for (int index : indexes) {
- elements.add(TestSource.element(tag, 0, index));
+ elements.add(TestUnboundedSource.element(tag, 0, index));
}
return elements;
}
-
- //
---------------------------------------------------------------------------------------------
- // the shared in memory UnboundedSource
- //
---------------------------------------------------------------------------------------------
-
- /**
- * Splits into one sub source per shard, each over {@code count / shards}
elements named {@code
- * <tag>-<shard>-<index>} with evenly spaced timestamps. Marks are not Java
serializable, they
- * record the position they finalize under {@code <tag>/<shard>}, readers
are counted per tag.
- */
- static final class TestSource extends UnboundedSource<String,
TestSource.Mark> {
- private static final long serialVersionUID = 1L;
-
- static final Coder<Mark> MARK_CODER = new MarkCoder();
-
- private static final ConcurrentMap<String, List<Integer>> FINALIZED = new
ConcurrentHashMap<>();
- private static final ConcurrentMap<String, AtomicInteger> CREATED = new
ConcurrentHashMap<>();
-
- private final String tag;
- private final int shard;
- private final int shards;
- private final int perShard;
-
- TestSource(String tag, int shards, int count) {
- this(tag, -1, shards, count / shards);
- }
-
- private TestSource(String tag, int shard, int shards, int perShard) {
- this.tag = tag;
- this.shard = shard;
- this.shards = shards;
- this.perShard = perShard;
- }
-
- static Set<String> elements(String tag, int shards, int count) {
- Set<String> elements = new HashSet<>();
- for (int shard = 0; shard < shards; shard++) {
- for (int index = 0; index < count / shards; index++) {
- elements.add(element(tag, shard, index));
- }
- }
- return elements;
- }
-
- static String element(String tag, int shard, int index) {
- return tag + "-" + shard + "-" + index;
- }
-
- static int shardOf(String element) {
- String head = element.substring(0, element.lastIndexOf('-'));
- return Integer.parseInt(head.substring(head.lastIndexOf('-') + 1));
- }
-
- static int indexOf(String element) {
- return Integer.parseInt(element.substring(element.lastIndexOf('-') + 1));
- }
-
- static List<Integer> finalized(String tag, int shard) {
- List<Integer> positions = FINALIZED.get(key(tag, shard));
- if (positions == null) {
- return Collections.emptyList();
- }
- synchronized (positions) {
- return new ArrayList<>(positions);
- }
- }
-
- static int created(String tag) {
- AtomicInteger created = CREATED.get(tag);
- return created == null ? 0 : created.get();
- }
-
- static void forget(String tag) {
- FINALIZED.keySet().removeIf(key -> key.startsWith(tag + "/"));
- CREATED.remove(tag);
- }
-
- private static String key(String tag, int shard) {
- return tag + "/" + shard;
- }
-
- @Override
- public List<TestSource> split(int desiredNumSplits, PipelineOptions
options) {
- if (shard >= 0) {
- return Collections.singletonList(this);
- }
- List<TestSource> splits = new ArrayList<>();
- for (int i = 0; i < shards; i++) {
- splits.add(new TestSource(tag, i, shards, perShard));
- }
- return splits;
- }
-
- @Override
- public UnboundedReader<String> createReader(PipelineOptions options,
@Nullable Mark mark) {
- if (shard < 0) {
- throw new IllegalStateException("split before reading");
- }
- CREATED.computeIfAbsent(tag, t -> new AtomicInteger()).incrementAndGet();
- return new Reader(this, mark == null ? 0 : mark.next);
- }
-
- @Override
- public Coder<Mark> getCheckpointMarkCoder() {
- return MARK_CODER;
- }
-
- @Override
- public Coder<String> getOutputCoder() {
- return StringUtf8Coder.of();
- }
-
- /** Position of the next element of a shard, deliberately not {@link
Serializable}. */
- static final class Mark implements UnboundedSource.CheckpointMark {
- private final String tag;
- private final int shard;
- final int next;
-
- Mark(String tag, int shard, int next) {
- this.tag = tag;
- this.shard = shard;
- this.next = next;
- }
-
- @Override
- public void finalizeCheckpoint() {
- FINALIZED
- .computeIfAbsent(key(tag, shard), k ->
Collections.synchronizedList(new ArrayList<>()))
- .add(next);
- }
- }
-
- private static final class MarkCoder extends CustomCoder<Mark> {
- private static final long serialVersionUID = 1L;
-
- @Override
- public void encode(Mark mark, OutputStream out) throws IOException {
- StringUtf8Coder.of().encode(mark.tag, out);
- VarIntCoder.of().encode(mark.shard, out);
- VarIntCoder.of().encode(mark.next, out);
- }
-
- @Override
- public Mark decode(InputStream in) throws IOException {
- return new Mark(
- StringUtf8Coder.of().decode(in),
- VarIntCoder.of().decode(in),
- VarIntCoder.of().decode(in));
- }
- }
-
- private static final class Reader extends UnboundedReader<String> {
- private final TestSource source;
- private int next;
- private int current = -1;
-
- Reader(TestSource source, int next) {
- this.source = source;
- this.next = next;
- }
-
- @Override
- public boolean start() {
- return advance();
- }
-
- @Override
- public boolean advance() {
- if (next < source.perShard) {
- current = next++;
- return true;
- }
- return false;
- }
-
- @Override
- public String getCurrent() throws NoSuchElementException {
- if (current < 0) {
- throw new NoSuchElementException();
- }
- return element(source.tag, source.shard, current);
- }
-
- @Override
- public Instant getCurrentTimestamp() throws NoSuchElementException {
- if (current < 0) {
- throw new NoSuchElementException();
- }
- return new Instant(
- BASE_MILLIS + (source.shard * source.perShard + current) *
INTERVAL_MILLIS);
- }
-
- @Override
- public Instant getWatermark() {
- return current < 0 ? BoundedWindow.TIMESTAMP_MIN_VALUE :
getCurrentTimestamp();
- }
-
- @Override
- public CheckpointMark getCheckpointMark() {
- return new Mark(source.tag, source.shard, next);
- }
-
- @Override
- public UnboundedSource<String, ?> getCurrentSource() {
- return source;
- }
-
- @Override
- public void close() {}
- }
- }
}
diff --git
a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/TestUnboundedSource.java
b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/TestUnboundedSource.java
new file mode 100644
index 00000000000..4a71dee5ac8
--- /dev/null
+++
b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/TestUnboundedSource.java
@@ -0,0 +1,252 @@
+/*
+ * 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.io.streaming;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.NoSuchElementException;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.atomic.AtomicInteger;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CustomCoder;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.coders.VarIntCoder;
+import org.apache.beam.sdk.io.UnboundedSource;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Instant;
+
+/**
+ * In memory unbounded source for tests. Splits into one sub source per shard.
Elements are strings
+ * formatted as tag, shard, index. Watermark freezes after exhaustion. Marks
record finalized
+ * positions under tag and shard.
+ */
+public final class TestUnboundedSource extends UnboundedSource<String,
TestUnboundedSource.Mark> {
+ private static final long serialVersionUID = 1L;
+
+ /** A modern timestamp with no rebase or DST subtleties. */
+ public static final long BASE_MILLIS = 1_700_000_000_000L;
+
+ public static final long INTERVAL_MILLIS = 1_000L;
+
+ public static final Coder<Mark> MARK_CODER = new MarkCoder();
+
+ private static final ConcurrentMap<String, List<Integer>> FINALIZED = new
ConcurrentHashMap<>();
+ private static final ConcurrentMap<String, AtomicInteger> CREATED = new
ConcurrentHashMap<>();
+
+ private final String tag;
+ private final int shard;
+ private final int shards;
+ private final int perShard;
+
+ public TestUnboundedSource(String tag, int shards, int count) {
+ this(tag, -1, shards, count / shards);
+ }
+
+ private TestUnboundedSource(String tag, int shard, int shards, int perShard)
{
+ this.tag = tag;
+ this.shard = shard;
+ this.shards = shards;
+ this.perShard = perShard;
+ }
+
+ public static Set<String> elements(String tag, int shards, int count) {
+ Set<String> elements = new HashSet<>();
+ for (int shard = 0; shard < shards; shard++) {
+ for (int index = 0; index < count / shards; index++) {
+ elements.add(element(tag, shard, index));
+ }
+ }
+ return elements;
+ }
+
+ public static String element(String tag, int shard, int index) {
+ return tag + "-" + shard + "-" + index;
+ }
+
+ public static int shardOf(String element) {
+ String head = element.substring(0, element.lastIndexOf('-'));
+ return Integer.parseInt(head.substring(head.lastIndexOf('-') + 1));
+ }
+
+ public static int indexOf(String element) {
+ return Integer.parseInt(element.substring(element.lastIndexOf('-') + 1));
+ }
+
+ public static List<Integer> finalized(String tag, int shard) {
+ List<Integer> positions = FINALIZED.get(key(tag, shard));
+ if (positions == null) {
+ return Collections.emptyList();
+ }
+ synchronized (positions) {
+ return new ArrayList<>(positions);
+ }
+ }
+
+ public static int created(String tag) {
+ AtomicInteger created = CREATED.get(tag);
+ return created == null ? 0 : created.get();
+ }
+
+ public static void forget(String tag) {
+ FINALIZED.keySet().removeIf(key -> key.startsWith(tag + "/"));
+ CREATED.remove(tag);
+ }
+
+ private static String key(String tag, int shard) {
+ return tag + "/" + shard;
+ }
+
+ @Override
+ public List<TestUnboundedSource> split(int desiredNumSplits, PipelineOptions
options) {
+ if (shard >= 0) {
+ return Collections.singletonList(this);
+ }
+ List<TestUnboundedSource> splits = new ArrayList<>();
+ for (int i = 0; i < shards; i++) {
+ splits.add(new TestUnboundedSource(tag, i, shards, perShard));
+ }
+ return splits;
+ }
+
+ @Override
+ public UnboundedReader<String> createReader(PipelineOptions options,
@Nullable Mark mark) {
+ if (shard < 0) {
+ throw new IllegalStateException("split before reading");
+ }
+ CREATED.computeIfAbsent(tag, t -> new AtomicInteger()).incrementAndGet();
+ return new Reader(this, mark == null ? 0 : mark.next);
+ }
+
+ @Override
+ public Coder<Mark> getCheckpointMarkCoder() {
+ return MARK_CODER;
+ }
+
+ @Override
+ public Coder<String> getOutputCoder() {
+ return StringUtf8Coder.of();
+ }
+
+ /** Position of the next element of a shard, not Java serializable. */
+ public static final class Mark implements UnboundedSource.CheckpointMark {
+ private final String tag;
+ private final int shard;
+ final int next;
+
+ public Mark(String tag, int shard, int next) {
+ this.tag = tag;
+ this.shard = shard;
+ this.next = next;
+ }
+
+ @Override
+ public void finalizeCheckpoint() {
+ FINALIZED
+ .computeIfAbsent(key(tag, shard), k ->
Collections.synchronizedList(new ArrayList<>()))
+ .add(next);
+ }
+ }
+
+ private static final class MarkCoder extends CustomCoder<Mark> {
+ private static final long serialVersionUID = 1L;
+
+ @Override
+ public void encode(Mark mark, OutputStream out) throws IOException {
+ StringUtf8Coder.of().encode(mark.tag, out);
+ VarIntCoder.of().encode(mark.shard, out);
+ VarIntCoder.of().encode(mark.next, out);
+ }
+
+ @Override
+ public Mark decode(InputStream in) throws IOException {
+ return new Mark(
+ StringUtf8Coder.of().decode(in),
+ VarIntCoder.of().decode(in),
+ VarIntCoder.of().decode(in));
+ }
+ }
+
+ private static final class Reader extends UnboundedReader<String> {
+ private final TestUnboundedSource source;
+ private int next;
+ private int current = -1;
+
+ Reader(TestUnboundedSource source, int next) {
+ this.source = source;
+ this.next = next;
+ }
+
+ @Override
+ public boolean start() {
+ return advance();
+ }
+
+ @Override
+ public boolean advance() {
+ if (next < source.perShard) {
+ current = next++;
+ return true;
+ }
+ return false;
+ }
+
+ @Override
+ public String getCurrent() throws NoSuchElementException {
+ if (current < 0) {
+ throw new NoSuchElementException();
+ }
+ return element(source.tag, source.shard, current);
+ }
+
+ @Override
+ public Instant getCurrentTimestamp() throws NoSuchElementException {
+ if (current < 0) {
+ throw new NoSuchElementException();
+ }
+ return new Instant(
+ BASE_MILLIS + (source.shard * source.perShard + current) *
INTERVAL_MILLIS);
+ }
+
+ @Override
+ public Instant getWatermark() {
+ return current < 0 ? BoundedWindow.TIMESTAMP_MIN_VALUE :
getCurrentTimestamp();
+ }
+
+ @Override
+ public CheckpointMark getCheckpointMark() {
+ return new Mark(source.tag, source.shard, next);
+ }
+
+ @Override
+ public UnboundedSource<String, ?> getCurrentSource() {
+ return source;
+ }
+
+ @Override
+ public void close() {}
+ }
+}
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
new file mode 100644
index 00000000000..4ad37f78f4a
--- /dev/null
+++
b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreamingTest.java
@@ -0,0 +1,142 @@
+/*
+ * 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;
+
+import static org.junit.Assert.assertThrows;
+
+import java.io.Serializable;
+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.runners.spark.structuredstreaming.translation.streaming.StreamingTestUtils;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.coders.VarIntCoder;
+import org.apache.beam.sdk.io.Read;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.transforms.Combine;
+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.FixedWindows;
+import org.apache.beam.sdk.transforms.windowing.Window;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
+import org.joda.time.Duration;
+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;
+
+/** Verifies that unsupported transforms in streaming mode fail at translation
time. */
+@RunWith(JUnit4.class)
+@Category(StreamingTest.class)
+public class PipelineTranslatorStreamingTest implements Serializable {
+
+ @ClassRule public static final SparkSessionRule SESSION = new
SparkSessionRule();
+ @Rule public transient TemporaryFolder temp = new TemporaryFolder();
+
+ private PCollection<KV<String, Integer>> kv(String tag) throws Exception {
+ SparkStructuredStreamingPipelineOptions o =
StreamingTestUtils.streamingOptions(temp);
+ return Pipeline.create(o)
+ .apply(Read.from(new TestUnboundedSource(tag, 1, 1)))
+ .apply(Window.into(FixedWindows.of(Duration.millis(1))))
+ .apply(ParDo.of(new ToKvFn()));
+ }
+
+ private static void assertUnsupported(Pipeline pipeline, String expected) {
+ Throwable thrown = assertThrows(Exception.class, () ->
StreamingTestUtils.run(pipeline));
+ for (Throwable t = thrown; t != null; t = t.getCause()) {
+ if (t instanceof UnsupportedOperationException &&
t.getMessage().contains(expected)) {
+ return;
+ }
+ }
+ 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);
+ Pipeline p = Pipeline.create(o);
+ p.apply(Create.of("rejected"));
+ assertUnsupported(p, "Impulse");
+ }
+
+ @Test
+ public void rejectsBoundedReadFromCreate() throws Exception {
+ SparkStructuredStreamingPipelineOptions o =
StreamingTestUtils.streamingOptions(temp);
+ Pipeline p = Pipeline.create(o);
+ p.apply(Create.of("rejected", "too"));
+ assertUnsupported(p, "Bounded Read");
+ }
+
+ @Test
+ public void rejectsStatefulParDo() throws Exception {
+ assertUnsupported(kv("s").apply(ParDo.of(new
StatefulDoFn())).getPipeline(), "Stateful ParDo");
+ }
+
+ @Test
+ public void rejectsAdditionalOutputs() throws Exception {
+ TupleTag<Integer> main = new TupleTag<Integer>() {};
+ TupleTag<Integer> other = new TupleTag<Integer>() {};
+ assertUnsupported(
+ kv("o")
+ .apply(ParDo.of(new StatelessDoFn()).withOutputTags(main,
TupleTagList.of(other)))
+ .getPipeline(),
+ "additional outputs");
+ }
+
+ private static final class ToKvFn extends DoFn<String, KV<String, Integer>> {
+ @ProcessElement
+ public void process(@Element String element, OutputReceiver<KV<String,
Integer>> out) {
+ out.output(KV.of(element, 1));
+ }
+ }
+
+ private static final class StatelessDoFn extends DoFn<KV<String, Integer>,
Integer> {
+ @ProcessElement
+ public void process() {}
+ }
+
+ private static final class StatefulDoFn extends DoFn<KV<String, Integer>,
Integer> {
+ @DoFn.StateId("state")
+ final StateSpec<?> spec = StateSpecs.value(VarIntCoder.of());
+
+ @ProcessElement
+ public void process() {}
+ }
+}
diff --git
a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatelessParDoStreamingTest.java
b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatelessParDoStreamingTest.java
new file mode 100644
index 00000000000..d41f5e77422
--- /dev/null
+++
b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatelessParDoStreamingTest.java
@@ -0,0 +1,110 @@
+/*
+ * 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 java.io.Serializable;
+import java.util.HashSet;
+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.io.Read;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.Flatten;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionList;
+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;
+
+/** Baseline streaming pipeline tests for stateless ParDo and Flatten. */
+@RunWith(JUnit4.class)
+@Category(StreamingTest.class)
+public class StatelessParDoStreamingTest implements Serializable {
+
+ @ClassRule public static final SparkSessionRule SESSION = new
SparkSessionRule();
+
+ @Rule public transient TemporaryFolder checkpointDir = new TemporaryFolder();
+
+ private static class PassThroughFn extends DoFn<String, String> {
+ @ProcessElement
+ public void process(@Element String element, OutputReceiver<String> out) {
+ out.output(element);
+ }
+ }
+
+ @Test
+ public void everyElementPassesThrough() throws Exception {
+ String tag = "stateless-pardo";
+ String collectorId = StreamingTestUtils.newCollectorId(tag);
+ StreamingTestUtils.clear(collectorId);
+
+ SparkStructuredStreamingPipelineOptions options =
+ StreamingTestUtils.streamingOptions(checkpointDir);
+ Pipeline pipeline = Pipeline.create(options);
+
+ pipeline
+ .apply("ReadUnbounded", Read.from(new TestUnboundedSource(tag, 1, 10)))
+ .apply("PassThrough", ParDo.of(new PassThroughFn()))
+ .apply("Collect", ParDo.of(new
StreamingTestUtils.CollectDoFn<>(collectorId)));
+
+ PipelineResult result = StreamingTestUtils.run(pipeline);
+
+ Set<String> collected = StreamingTestUtils.collected(collectorId);
+ Set<String> expected = TestUnboundedSource.elements(tag, 1, 10);
+ assertEquals("pipeline state=" + result.getState(), expected, collected);
+ }
+
+ @Test
+ public void flattenCombinesMultipleUnboundedSources() throws Exception {
+ String tagA = "flatten-a";
+ String tagB = "flatten-b";
+ String collectorId = StreamingTestUtils.newCollectorId("flatten");
+ StreamingTestUtils.clear(collectorId);
+
+ SparkStructuredStreamingPipelineOptions options =
+ StreamingTestUtils.streamingOptions(checkpointDir);
+ Pipeline pipeline = Pipeline.create(options);
+
+ PCollection<String> a = pipeline.apply("ReadA", Read.from(new
TestUnboundedSource(tagA, 1, 5)));
+ PCollection<String> b = pipeline.apply("ReadB", Read.from(new
TestUnboundedSource(tagB, 1, 5)));
+
+ PCollectionList.of(a)
+ .and(b)
+ .apply("Flatten", Flatten.pCollections())
+ .apply("Collect", ParDo.of(new
StreamingTestUtils.CollectDoFn<>(collectorId)));
+
+ PipelineResult result = StreamingTestUtils.run(pipeline);
+
+ Set<String> collected = StreamingTestUtils.collected(collectorId);
+ Set<String> expected = new HashSet<>();
+ expected.addAll(TestUnboundedSource.elements(tagA, 1, 5));
+ expected.addAll(TestUnboundedSource.elements(tagB, 1, 5));
+ assertEquals("pipeline state=" + result.getState(), expected, collected);
+ }
+}
diff --git
a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingCheckpointRestartTest.java
b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingCheckpointRestartTest.java
new file mode 100644
index 00000000000..c7ad2fa6295
--- /dev/null
+++
b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingCheckpointRestartTest.java
@@ -0,0 +1,118 @@
+/*
+ * 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.File;
+import java.io.Serializable;
+import java.util.Collections;
+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.BeamReaderCache;
+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.io.Read;
+import org.apache.beam.sdk.transforms.ParDo;
+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;
+
+/**
+ * Verifies restarted pipelines resume from durable checkpoint marks.
+ *
+ * <p>Two pipelines run sequentially against the same checkpoint directory.
Wiping the reader cache
+ * in between forces the second run to restore from the durable marks.
+ */
+@RunWith(JUnit4.class)
+@Category(StreamingTest.class)
+public class StreamingCheckpointRestartTest implements Serializable {
+
+ @ClassRule public static final SparkSessionRule SESSION = new
SparkSessionRule();
+
+ @Rule public transient TemporaryFolder checkpointDir = new TemporaryFolder();
+
+ private static final int ELEMENT_COUNT = 10;
+ private static final String TAG = "checkpoint-restart";
+
+ @After
+ public void tearDown() {
+ TestUnboundedSource.forget(TAG);
+ }
+
+ @Test
+ public void restartedPipelineResumesFromDurableCheckpointMarks() throws
Exception {
+ String checkpointPath =
checkpointDir.newFolder("checkpoint").getAbsolutePath();
+
+ String collectorA =
StreamingTestUtils.newCollectorId("checkpoint-restart-a");
+ String collectorB =
StreamingTestUtils.newCollectorId("checkpoint-restart-b");
+ StreamingTestUtils.clear(collectorA);
+ StreamingTestUtils.clear(collectorB);
+
+ // First run reads all elements from a fresh checkpoint directory.
+ runPipeline(checkpointPath, collectorA, TAG);
+
+ Set<String> collectedA = StreamingTestUtils.collected(collectorA);
+ assertEquals(
+ "first run must read every element",
+ TestUnboundedSource.elements(TAG, 1, ELEMENT_COUNT),
+ collectedA);
+
+ // The source checkpoint lives under the location the translator handed to
Spark.
+ File sourceRoot = new File(new File(checkpointPath, "0"), "sources/0");
+ assertTrue("expected source checkpoint directory " + sourceRoot,
sourceRoot.isDirectory());
+
+ int createdBeforeSecondRun = TestUnboundedSource.created(TAG);
+
+ // Invalidate cached readers and marks to force restore from durable files.
+ BeamReaderCache.invalidateAll();
+
+ // Second run resumes against the same checkpoint directory.
+ runPipeline(checkpointPath, collectorB, TAG);
+
+ Set<String> collectedB = StreamingTestUtils.collected(collectorB);
+ assertTrue(
+ "second run must not re-emit elements the first run committed",
+ Collections.disjoint(collectedA, collectedB));
+ assertTrue(
+ "readers must be recreated during the second run",
+ TestUnboundedSource.created(TAG) > createdBeforeSecondRun);
+ }
+
+ private PipelineResult runPipeline(String checkpointPath, String
collectorId, String tag) {
+ SparkStructuredStreamingPipelineOptions options =
+ StreamingTestUtils.streamingOptions(checkpointPath);
+ // One element per split per micro batch to spread elements across batches.
+ options.setMaxRecordsPerBatch(1L);
+
+ Pipeline pipeline = Pipeline.create(options);
+ pipeline
+ .apply("ReadUnbounded", Read.from(new TestUnboundedSource(tag, 1,
ELEMENT_COUNT)))
+ .apply("Collect", ParDo.of(new
StreamingTestUtils.CollectDoFn<>(collectorId)));
+ return StreamingTestUtils.run(pipeline);
+ }
+}
diff --git
a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java
b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java
new file mode 100644
index 00000000000..76462c38fe1
--- /dev/null
+++
b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java
@@ -0,0 +1,171 @@
+/*
+ * 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.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+import java.io.Serializable;
+import java.util.Set;
+import java.util.function.IntPredicate;
+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.io.Read;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.ParDo;
+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;
+
+/**
+ * State transitions for a streaming pipeline.
+ *
+ * <p>Pipelines observe RUNNING, DONE once idle, CANCELLED on cancel, and
FAILED on query failure.
+ */
+@RunWith(JUnit4.class)
+@Category(StreamingTest.class)
+public class StreamingPipelineLifecycleTest implements Serializable {
+
+ /** Session shared across tests. */
+ @ClassRule public static final SparkSessionRule SESSION = new
SparkSessionRule();
+
+ @Rule public transient TemporaryFolder checkpointDir = new TemporaryFolder();
+
+ /** Upper bound for a query to start or stop. */
+ private static final long QUERY_START_TIMEOUT_MILLIS = 60_000L;
+
+ /** Polls the active query count of the shared session until {@code
condition} holds. */
+ private static void awaitActiveQueries(IntPredicate condition, String
failure)
+ throws InterruptedException {
+ long deadline = System.currentTimeMillis() + QUERY_START_TIMEOUT_MILLIS;
+ while (!condition.test(SESSION.getSession().streams().active().length)) {
+ assertTrue(
+ failure + " within " + QUERY_START_TIMEOUT_MILLIS + " ms",
+ System.currentTimeMillis() < deadline);
+ Thread.sleep(50L);
+ }
+ }
+
+ @Test
+ public void idlePipelineGoesFromRunningToDoneOnceIdle() throws Exception {
+ String tag = "lifecycle-done";
+ String collectorId = StreamingTestUtils.newCollectorId(tag);
+ StreamingTestUtils.clear(collectorId);
+
+ SparkStructuredStreamingPipelineOptions options =
+ StreamingTestUtils.streamingOptions(checkpointDir);
+ Pipeline pipeline = Pipeline.create(options);
+
+ pipeline
+ .apply("ReadUnbounded", Read.from(new TestUnboundedSource(tag, 1, 10)))
+ .apply("Collect", ParDo.of(new
StreamingTestUtils.CollectDoFn<>(collectorId)));
+
+ PipelineResult result = pipeline.run();
+ assertEquals(PipelineResult.State.RUNNING, result.getState());
+
+ PipelineResult.State finalState =
StreamingTestUtils.waitUntilFinish(result);
+ assertEquals(PipelineResult.State.DONE, finalState);
+ assertEquals(PipelineResult.State.DONE, result.getState());
+
+ Set<String> collected = StreamingTestUtils.collected(collectorId);
+ assertEquals(TestUnboundedSource.elements(tag, 1, 10), collected);
+ }
+
+ @Test
+ public void cancelStopsTheQueryAndReportsCancelled() throws Exception {
+ String tag = "lifecycle-cancel";
+ String collectorId = StreamingTestUtils.newCollectorId(tag);
+ StreamingTestUtils.clear(collectorId);
+
+ SparkStructuredStreamingPipelineOptions options =
+ StreamingTestUtils.streamingOptions(checkpointDir);
+ // Idle stop disabled so the query stops only from explicit cancel.
+ options.setStreamingStopAfterIdleBatches(-1);
+ Pipeline pipeline = Pipeline.create(options);
+
+ pipeline
+ .apply("ReadUnbounded", Read.from(new TestUnboundedSource(tag, 1, 10)))
+ .apply("Collect", ParDo.of(new
StreamingTestUtils.CollectDoFn<>(collectorId)));
+
+ PipelineResult result = pipeline.run();
+ assertEquals(PipelineResult.State.RUNNING, result.getState());
+
+ awaitActiveQueries(count -> count > 0, "no streaming query started");
+
+ PipelineResult.State cancelledState = result.cancel();
+ assertEquals(PipelineResult.State.CANCELLED, cancelledState);
+ assertEquals(PipelineResult.State.CANCELLED, result.getState());
+
+ awaitActiveQueries(count -> count == 0, "the streaming query did not stop
after cancel");
+ }
+
+ /** A failure in any leaf query surfaces through waitUntilFinish. */
+ @Test
+ public void failingLeafQueryFailsThePipelineAndStopsHealthySibling() throws
Exception {
+ String healthyTag = "lifecycle-healthy";
+ String poisonTag = "lifecycle-poison";
+ String collectorId = StreamingTestUtils.newCollectorId(healthyTag);
+ StreamingTestUtils.clear(collectorId);
+
+ SparkStructuredStreamingPipelineOptions options =
+ StreamingTestUtils.streamingOptions(checkpointDir);
+ // Idle stop disabled so the healthy query stops only when the sibling
failure stops it.
+ options.setStreamingStopAfterIdleBatches(-1);
+ Pipeline pipeline = Pipeline.create(options);
+
+ pipeline
+ .apply("ReadHealthy", Read.from(new TestUnboundedSource(healthyTag, 1,
10)))
+ .apply("Collect", ParDo.of(new
StreamingTestUtils.CollectDoFn<>(collectorId)));
+ pipeline
+ .apply("ReadPoisoned", Read.from(new TestUnboundedSource(poisonTag, 1,
10)))
+ .apply("Throw", ParDo.of(new ThrowOnElementDoFn(5)));
+
+ PipelineResult result = pipeline.run();
+
+ assertThrows(RuntimeException.class, () ->
StreamingTestUtils.waitUntilFinish(result));
+ assertEquals(PipelineResult.State.FAILED, result.getState());
+
+ awaitActiveQueries(count -> count == 0, "a sibling query did not stop
after the failure");
+ }
+
+ /** Throws on one specific element index, passes every other element
through. */
+ private static final class ThrowOnElementDoFn extends DoFn<String, String> {
+ private final int poisonIndex;
+
+ ThrowOnElementDoFn(int poisonIndex) {
+ this.poisonIndex = poisonIndex;
+ }
+
+ @ProcessElement
+ public void processElement(@Element String element, OutputReceiver<String>
out) {
+ if (TestUnboundedSource.indexOf(element) == poisonIndex) {
+ throw new IllegalStateException("poison index " + poisonIndex);
+ }
+ out.output(element);
+ }
+ }
+}
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
new file mode 100644
index 00000000000..3d784115415
--- /dev/null
+++
b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTestUtils.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;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.ConcurrentHashMap;
+import
org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions;
+import
org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingRunner;
+import org.apache.beam.sdk.Pipeline;
+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.joda.time.Duration;
+import org.junit.rules.TemporaryFolder;
+
+/**
+ * Shared test utilities for the Spark 4 streaming translators.
+ *
+ * <p>Collectors are static so they work in local mode only.
+ *
+ * <p>The {@link #run} helper bounds each query at five minutes and cancels on
expiry.
+ *
+ * <p>Tests poll with deadlines instead of {@code @Test(timeout)}, JUnit runs
a timed test in a
+ * separate thread group and Spark's static thread pools inherit it.
+ */
+public final class StreamingTestUtils {
+
+ private StreamingTestUtils() {}
+
+ /** Driver side, per collector id accumulation of every element a {@link
CollectDoFn} saw. */
+ private static final Map<String, List<Object>> COLLECTORS = new
ConcurrentHashMap<>();
+
+ /**
+ * Appends every element to a static collector named {@code collectorId},
then passes it through.
+ * Safe to use concurrently. Works in Spark local mode only.
+ */
+ public static final class CollectDoFn<T> extends DoFn<T, T> {
+ private final String collectorId;
+
+ public CollectDoFn(String collectorId) {
+ this.collectorId = Preconditions.checkNotNull(collectorId);
+ }
+
+ @ProcessElement
+ public void processElement(@Element T element, OutputReceiver<T> out) {
+ append(collectorId, element);
+ out.output(element);
+ }
+ }
+
+ private static void append(String collectorId, Object value) {
+ COLLECTORS
+ .computeIfAbsent(collectorId, unused ->
Collections.synchronizedList(new ArrayList<>()))
+ .add(value);
+ }
+
+ /** Returns a snapshot of everything collected so far under {@code
collectorId}. */
+ @SuppressWarnings("unchecked")
+ public static <T> Set<T> collected(String collectorId) {
+ List<Object> values = COLLECTORS.get(collectorId);
+ if (values == null) {
+ return Collections.emptySet();
+ }
+ synchronized (values) {
+ return (Set<T>) new HashSet<>(values);
+ }
+ }
+
+ /** Discards everything collected so far under {@code collectorId}. */
+ public static void clear(String collectorId) {
+ COLLECTORS.remove(collectorId);
+ }
+
+ /** Collector id that will not collide with other tests or runs. */
+ public static String newCollectorId(String prefix) {
+ return prefix + "-" + UUID.randomUUID();
+ }
+
+ /** Upper bound on the wall clock time one streaming pipeline may take. */
+ public static final Duration FINISH_TIMEOUT = Duration.standardMinutes(5);
+
+ /**
+ * Runs {@code pipeline} and returns on a terminal state. Cancels the
pipeline and fails after
+ * {@link #FINISH_TIMEOUT}.
+ */
+ public static PipelineResult run(Pipeline pipeline) {
+ PipelineResult result = pipeline.run();
+ waitUntilFinish(result);
+ return result;
+ }
+
+ /**
+ * Waits at most {@link #FINISH_TIMEOUT} for {@code result} to reach a
terminal state. Fails and
+ * cancels if still running at the deadline.
+ */
+ public static PipelineResult.State waitUntilFinish(PipelineResult result) {
+ PipelineResult.State state = result.waitUntilFinish(FINISH_TIMEOUT);
+ if (state == null || !state.isTerminal()) {
+ try {
+ result.cancel();
+ } catch (IOException | RuntimeException e) {
+ // Best effort cancellation.
+ }
+ throw new AssertionError(
+ "pipeline did not finish within " + FINISH_TIMEOUT + ", last state "
+ state);
+ }
+ return state;
+ }
+
+ /**
+ * Streaming options for tests: runner on the active session, streaming
mode, stop after 3 idle
+ * batches, 200 ms trigger, the given checkpoint directory. Test mode is
off, {@code run()}
+ * returns at once and tests wait through {@link #run} or {@link
#waitUntilFinish}.
+ */
+ public static SparkStructuredStreamingPipelineOptions streamingOptions(
+ TemporaryFolder checkpointDir) throws IOException {
+ return
streamingOptions(checkpointDir.newFolder("checkpoint").getAbsolutePath());
+ }
+
+ /** Same as {@link #streamingOptions(TemporaryFolder)} with an explicit
checkpoint path. */
+ public static SparkStructuredStreamingPipelineOptions
streamingOptions(String checkpointPath) {
+ SparkStructuredStreamingPipelineOptions options =
+
PipelineOptionsFactory.as(SparkStructuredStreamingPipelineOptions.class);
+ options.setRunner(SparkStructuredStreamingRunner.class);
+ options.setUseActiveSparkSession(true);
+ options.setTestMode(false);
+ options.setStreaming(true);
+ options.setStreamingStopAfterIdleBatches(3);
+ options.setMaxBatchDurationMillis(200);
+ options.setCheckpointDir(checkpointPath);
+ return options;
+ }
+}
diff --git
a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/ParDoTranslatorBatch.java
b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/ParDoTranslatorBatch.java
index 14058a73733..b454fc1a34e 100644
---
a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/ParDoTranslatorBatch.java
+++
b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/ParDoTranslatorBatch.java
@@ -92,7 +92,7 @@ class ParDoTranslatorBatch<InputT, OutputT>
doFn);
// Stateful, timer using and time sorted DoFns are routed to
StatefulParDoTranslatorBatch by
- // PipelineTranslatorBatch#getTransformTranslator. Reaching here with one
means dispatch is
+ // PipelineTranslatorCommon#getTransformTranslator. Reaching here with one
means dispatch is
// broken, not that the feature is unsupported.
checkState(
!StatefulParDoTranslatorBatch.appliesTo(transform),
diff --git
a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/PipelineTranslatorBatch.java
b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/PipelineTranslatorBatch.java
index ba7cbb0fa03..df8977063cd 100644
---
a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/PipelineTranslatorBatch.java
+++
b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/PipelineTranslatorBatch.java
@@ -17,95 +17,8 @@
*/
package org.apache.beam.runners.spark.structuredstreaming.translation.batch;
-import java.util.HashMap;
-import java.util.Map;
-import
org.apache.beam.runners.spark.structuredstreaming.translation.PipelineTranslator;
-import
org.apache.beam.runners.spark.structuredstreaming.translation.TransformTranslator;
-import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.annotations.Internal;
-import org.apache.beam.sdk.transforms.Combine;
-import org.apache.beam.sdk.transforms.Flatten;
-import org.apache.beam.sdk.transforms.GroupByKey;
-import org.apache.beam.sdk.transforms.Impulse;
-import org.apache.beam.sdk.transforms.PTransform;
-import org.apache.beam.sdk.transforms.ParDo;
-import org.apache.beam.sdk.transforms.Reshuffle;
-import org.apache.beam.sdk.transforms.windowing.Window;
-import org.apache.beam.sdk.util.construction.SplittableParDo;
-import org.apache.beam.sdk.values.PInput;
-import org.apache.beam.sdk.values.POutput;
-import org.checkerframework.checker.nullness.qual.Nullable;
-/**
- * {@link PipelineTranslator} for executing a {@link Pipeline} in Spark in
batch mode. This contains
- * only the components specific to batch: registry of batch {@link
TransformTranslator} and registry
- * lookup code.
- */
+/** Translator for batch pipelines, the registry and the lookup live in the
common base. */
@Internal
-public class PipelineTranslatorBatch extends PipelineTranslator {
-
- //
--------------------------------------------------------------------------------------------
- // Transform Translator Registry
- //
--------------------------------------------------------------------------------------------
-
- @SuppressWarnings("rawtypes")
- private static final Map<Class<? extends PTransform>, TransformTranslator>
TRANSFORM_TRANSLATORS =
- new HashMap<>();
-
- // TODO the ability to have more than one TransformTranslator per URN
- // that could be dynamically chosen by a predicated that evaluates based on
PCollection
- // obtainable though node.getInputs.getValue()
- // See
- //
https://github.com/seznam/euphoria/blob/master/euphoria-spark/src/main/java/cz/seznam/euphoria/spark/SparkFlowTranslator.java#L83
- // And
- //
https://github.com/seznam/euphoria/blob/master/euphoria-spark/src/main/java/cz/seznam/euphoria/spark/SparkFlowTranslator.java#L106
-
- static {
- TRANSFORM_TRANSLATORS.put(Impulse.class, new ImpulseTranslatorBatch());
- TRANSFORM_TRANSLATORS.put(Combine.PerKey.class, new
CombinePerKeyTranslatorBatch<>());
- TRANSFORM_TRANSLATORS.put(Combine.Globally.class, new
CombineGloballyTranslatorBatch<>());
- TRANSFORM_TRANSLATORS.put(
- Combine.GroupedValues.class, new
CombineGroupedValuesTranslatorBatch<>());
- TRANSFORM_TRANSLATORS.put(GroupByKey.class, new
GroupByKeyTranslatorBatch<>());
-
- TRANSFORM_TRANSLATORS.put(Reshuffle.class, new
ReshuffleTranslatorBatch<>());
- TRANSFORM_TRANSLATORS.put(
- Reshuffle.ViaRandomKey.class, new
ReshuffleTranslatorBatch.ViaRandomKey<>());
-
- TRANSFORM_TRANSLATORS.put(Flatten.PCollections.class, new
FlattenTranslatorBatch<>());
-
- TRANSFORM_TRANSLATORS.put(Window.Assign.class, new
WindowAssignTranslatorBatch<>());
-
- TRANSFORM_TRANSLATORS.put(ParDo.MultiOutput.class, new
ParDoTranslatorBatch<>());
-
- TRANSFORM_TRANSLATORS.put(
- SplittableParDo.PrimitiveBoundedRead.class, new
ReadSourceTranslatorBatch<>());
- }
-
- /**
- * Translators that shadow the {@link #TRANSFORM_TRANSLATORS} entry for
their transform class when
- * a predicate matches, so that a single transform class can be translated
in more than one way
- * depending on the transform instance.
- *
- * <p>Currently only {@link ParDo.MultiOutput} needs this, to route stateful
and time sorted
- * {@link org.apache.beam.sdk.transforms.DoFn DoFns} away from {@link
ParDoTranslatorBatch}.
- */
- @SuppressWarnings("rawtypes")
- private static final TransformTranslator STATEFUL_PARDO_TRANSLATOR =
- new StatefulParDoTranslatorBatch<>();
-
- /** Returns a {@link TransformTranslator} for the given {@link PTransform}
if known. */
- @Override
- @Nullable
- protected <InT extends PInput, OutT extends POutput, TransformT extends
PTransform<InT, OutT>>
- TransformTranslator<InT, OutT, TransformT>
getTransformTranslator(TransformT transform) {
- // Resolved ahead of the class keyed registry: ParDo.MultiOutput maps to a
different translator
- // depending on the DoFn signature, which a lookup by transform class
alone cannot express. This
- // is the predicated dispatch of the TODO above, limited to the single
transform needing it.
- if (transform instanceof ParDo.MultiOutput
- && StatefulParDoTranslatorBatch.appliesTo((ParDo.MultiOutput<?, ?>)
transform)) {
- return STATEFUL_PARDO_TRANSLATOR;
- }
- return TRANSFORM_TRANSLATORS.get(transform.getClass());
- }
-}
+public class PipelineTranslatorBatch extends PipelineTranslatorCommon {}
diff --git
a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/PipelineTranslatorBatch.java
b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/PipelineTranslatorCommon.java
similarity index 94%
copy from
runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/PipelineTranslatorBatch.java
copy to
runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/PipelineTranslatorCommon.java
index ba7cbb0fa03..899e0443a94 100644
---
a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/PipelineTranslatorBatch.java
+++
b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/PipelineTranslatorCommon.java
@@ -21,7 +21,6 @@ import java.util.HashMap;
import java.util.Map;
import
org.apache.beam.runners.spark.structuredstreaming.translation.PipelineTranslator;
import
org.apache.beam.runners.spark.structuredstreaming.translation.TransformTranslator;
-import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.annotations.Internal;
import org.apache.beam.sdk.transforms.Combine;
import org.apache.beam.sdk.transforms.Flatten;
@@ -37,12 +36,11 @@ import org.apache.beam.sdk.values.POutput;
import org.checkerframework.checker.nullness.qual.Nullable;
/**
- * {@link PipelineTranslator} for executing a {@link Pipeline} in Spark in
batch mode. This contains
- * only the components specific to batch: registry of batch {@link
TransformTranslator} and registry
- * lookup code.
+ * Holds the translator registry shared by the batch and the streaming
translator, and the registry
+ * lookup.
*/
@Internal
-public class PipelineTranslatorBatch extends PipelineTranslator {
+public class PipelineTranslatorCommon extends PipelineTranslator {
//
--------------------------------------------------------------------------------------------
// Transform Translator Registry
diff --git
a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/StatefulParDoTranslatorBatch.java
b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/StatefulParDoTranslatorBatch.java
index 75d84b630bb..5f060b6aef4 100644
---
a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/StatefulParDoTranslatorBatch.java
+++
b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/StatefulParDoTranslatorBatch.java
@@ -55,7 +55,7 @@ import scala.Tuple2;
/**
* Translator for a stateful {@link ParDo.MultiOutput}, or one requiring time
sorted input.
*
- * <p>Selected by {@link PipelineTranslatorBatch} in place of {@link
ParDoTranslatorBatch} when the
+ * <p>Selected by {@link PipelineTranslatorCommon} in place of {@link
ParDoTranslatorBatch} when the
* {@link DoFn} uses state, uses timers, or is annotated with {@link
DoFn.RequiresTimeSortedInput};
* see {@link #appliesTo}.
*
