From a2ba7f57912cd9448cd945374369fded9055dc95 Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Thu, 10 Sep 2026 11:10:49 +0000 Subject: [PATCH 1/6] [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. StreamingEvaluationContext runs one noop sink query per leaf, checkpoints under checkpointDir/, stops siblings when a query fails and stops a query after streamingStopAfterIdleBatches triggers without input. The test source of BeamMicroBatchSourceTest moves to TestUnboundedSource so the translator tests share it. --- .../PipelineTranslatorFactory.java | 35 +++ .../PipelineTranslatorStreaming.java | 106 +++++++ .../StreamingEvaluationContext.java | 258 +++++++++++++++++ .../streaming/ReadUnboundedTranslator.java | 97 +++++++ .../streaming/BeamMicroBatchSourceTest.java | 269 ++---------------- .../io/streaming/TestUnboundedSource.java | 252 ++++++++++++++++ .../PipelineTranslatorStreamingTest.java | 142 +++++++++ .../StatelessParDoStreamingTest.java | 118 ++++++++ .../StreamingCheckpointRestartTest.java | 120 ++++++++ .../StreamingPipelineLifecycleTest.java | 193 +++++++++++++ .../streaming/StreamingTestUtils.java | 154 ++++++++++ 11 files changed, 1499 insertions(+), 245 deletions(-) create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorFactory.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/StreamingEvaluationContext.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ReadUnboundedTranslator.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/TestUnboundedSource.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreamingTest.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatelessParDoStreamingTest.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingCheckpointRestartTest.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTestUtils.java 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 000000000000..9f1c15a47773 --- /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 000000000000..d0ab3c25b615 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java @@ -0,0 +1,106 @@ +/* + * 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.PipelineTranslatorBatch; +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 batch translator 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 PipelineTranslatorBatch { + + /** Returns a {@link TransformTranslator} for the given {@link PTransform} if known. */ + @Override + @SuppressWarnings({"rawtypes", "unchecked"}) + @Nullable + protected > + TransformTranslator getTransformTranslator(TransformT transform) { + + if (transform instanceof SplittableParDo.PrimitiveUnboundedRead) { + return (TransformTranslator) new ReadUnboundedTranslator<>(); + } + + if (transform instanceof SplittableParDo.PrimitiveBoundedRead) { + throw unsupported( + "Bounded Read (Read.from(BoundedSource), Create with two or more elements)"); + } + + if (transform instanceof Impulse) { + throw unsupported("Impulse (Create with fewer than two elements, PAssert)"); + } + + if (transform instanceof GroupByKey) { + throw unsupported("GroupByKey"); + } + + if (transform instanceof Combine.PerKey) { + throw unsupported("Combine.perKey"); + } + + if (transform instanceof ParDo.MultiOutput) { + ParDo.MultiOutput parDo = (ParDo.MultiOutput) transform; + DoFnSignature signature = DoFnSignatures.signatureForDoFn(parDo.getFn()); + if (signature.usesState() || signature.usesTimers()) { + throw unsupported("Stateful ParDo (" + signature.fnClass().getName() + ")"); + } + if (!parDo.getSideInputs().isEmpty()) { + throw unsupported("ParDo with side inputs (" + signature.fnClass().getName() + ")"); + } + if (!parDo.getAdditionalOutputTags().getAll().isEmpty()) { + throw unsupported("ParDo with additional outputs (" + signature.fnClass().getName() + ")"); + } + } + + return super.getTransformTranslator(transform); + } + + private static UnsupportedOperationException unsupported(String what) { + return new UnsupportedOperationException( + what + + " is not supported by the Spark 4 streaming runner yet, see" + + " https://github.com/apache/beam/issues/36841"); + } + + @Override + protected EvaluationContext createEvaluationContext( + Collection> 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 000000000000..1edf3ea748cd --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/StreamingEvaluationContext.java @@ -0,0 +1,258 @@ +/* + * 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.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. + * + *

Leaf {@code i} checkpoints under {@code /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 queries = new ArrayList<>(); + private boolean stopped = false; + + StreamingEvaluationContext( + Collection> 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 toAwait; + synchronized (lock) { + toAwait = new ArrayList<>(queries); + } + awaitTermination(toAwait); + } 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 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", checkpointBaseDir + "/" + leafIndex) + .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 toAwait) { + List active = new ArrayList<>(toAwait); + while (!active.isEmpty()) { + Iterator 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( + "Error while stopping streaming query {}: {}", + query.id(), + String.valueOf(e.getMessage())); + } + } + + 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 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 000000000000..fb1691108b71 --- /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}. + * + *

Elements arrive in the global window with the record timestamp. Downstream windowing requires + * an explicit {@code Window.Assign}. + */ +public class ReadUnboundedTranslator + extends TransformTranslator, SplittableParDo.PrimitiveUnboundedRead> { + + public ReadUnboundedTranslator() { + super(0.05f); + } + + @Override + protected void translate(SplittableParDo.PrimitiveUnboundedRead transform, Context cxt) { + PCollection output = cxt.getOutput(); + UnboundedSource source = transform.getSource(); + Coder elementCoder = output.getCoder(); + + WindowedValues.FullWindowedValueCoder payloadCoder = + WindowedValues.getFullCoder(elementCoder, GlobalWindow.Coder.INSTANCE); + + SparkStructuredStreamingPipelineOptions options = + cxt.getOptions().as(SparkStructuredStreamingPipelineOptions.class); + + Dataset rows = + UnboundedSourceDataset.of( + cxt.getSparkSession(), + source, + payloadCoder, + options, + cxt.getCurrentTransform().getFullName()); + + Encoder> encoder = + cxt.windowedEncoder(elementCoder, GlobalWindow.Coder.INSTANCE); + + Dataset> 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 implements MapFunction> { + private final Coder> coder; + + DecodePayload(Coder> coder) { + this.coder = coder; + } + + @Override + public WindowedValue 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 4616c7563a20..3b2e36bf5ef7 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 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.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.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.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> 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 optionsBroadcast; @@ -164,7 +150,7 @@ public void setUp() { 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 void testReadsElementsFromUnboundedSource() throws Exception { 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 void testMaxRecordsPerBatchIsSharedAcrossSplits() throws Exception { List 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 void testRestartResumesFromCommittedMark() throws Exception { BeamReaderCache.invalidateAll(); List firstValues = values(batches(first)); - Set all = TestSource.elements(tag, shards, count); + Set all = TestUnboundedSource.elements(tag, shards, count); query = start(rows(shards, count, limited(limit, 1_000L)), second, checkpointDir); try { await( @@ -377,8 +366,8 @@ public void testRestartResumesFromCommittedMark() throws Exception { 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 void testMarksAreFinalizedOnlyAfterSparkCommit() throws Exception { int finalizations = 0; for (int shard = 0; shard < 2; shard++) { int committed = committedPosition(checkpointDir, shard); - List finalized = TestSource.finalized(tag, shard); + List finalized = TestUnboundedSource.finalized(tag, shard); assertTrue( "shard " + shard + " finalized " + finalized + " beyond committed " + committed, finalized.isEmpty() || Collections.max(finalized) <= committed); @@ -460,7 +449,7 @@ public void testStoppedQueryFinalizesLastCommittedMarks() throws Exception { for (int shard = 0; shard < 2; shard++) { int committed = committedPosition(checkpointDir, shard); - List finalized = TestSource.finalized(tag, shard); + List finalized = TestUnboundedSource.finalized(tag, shard); assertTrue( "shard " + shard + " finalized " + finalized + ", committed " + committed, finalized.contains(committed) && Collections.max(finalized) == committed); @@ -473,8 +462,8 @@ public void testRetriedBatchRestartsFromDurableMark() throws Exception { 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 void testMissingMarkThrows() throws Exception { 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 void testRetryAfterFailedMarkWriteRecreatesReader() throws Exception { 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 @@ private Dataset rows( 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 @@ private static List nonEmptySizes(List> batches) { private static Set shardsOf(List values) { Set shards = new HashSet<>(); for (String value : values) { - shards.add(TestSource.shardOf(value)); + shards.add(TestUnboundedSource.shardOf(value)); } return shards; } @@ -759,7 +748,7 @@ private static int committedPosition(File checkpointDir, int shard) throws IOExc 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 @@ private static int committedPosition(File checkpointDir, int shard) throws IOExc /** Split 0 of a single shard source of 100 elements from epoch {@code start} to {@code end}. */ private BeamInputPartition 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 @@ private static List drainUntilFailure( private List shardZero(int... indexes) { List 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 - * --} with evenly spaced timestamps. Marks are not Java serializable, they - * record the position they finalize under {@code /}, readers are counted per tag. - */ - static final class TestSource extends UnboundedSource { - private static final long serialVersionUID = 1L; - - static final Coder MARK_CODER = new MarkCoder(); - - private static final ConcurrentMap> FINALIZED = new ConcurrentHashMap<>(); - private static final ConcurrentMap 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 elements(String tag, int shards, int count) { - Set 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 finalized(String tag, int shard) { - List 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 split(int desiredNumSplits, PipelineOptions options) { - if (shard >= 0) { - return Collections.singletonList(this); - } - List splits = new ArrayList<>(); - for (int i = 0; i < shards; i++) { - splits.add(new TestSource(tag, i, shards, perShard)); - } - return splits; - } - - @Override - public UnboundedReader 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 getCheckpointMarkCoder() { - return MARK_CODER; - } - - @Override - public Coder 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 { - 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 { - 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 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 000000000000..4a71dee5ac81 --- /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 { + 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_CODER = new MarkCoder(); + + private static final ConcurrentMap> FINALIZED = new ConcurrentHashMap<>(); + private static final ConcurrentMap 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 elements(String tag, int shards, int count) { + Set 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 finalized(String tag, int shard) { + List 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 split(int desiredNumSplits, PipelineOptions options) { + if (shard >= 0) { + return Collections.singletonList(this); + } + List splits = new ArrayList<>(); + for (int i = 0; i < shards; i++) { + splits.add(new TestUnboundedSource(tag, i, shards, perShard)); + } + return splits; + } + + @Override + public UnboundedReader 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 getCheckpointMarkCoder() { + return MARK_CODER; + } + + @Override + public Coder 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 { + 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 { + 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 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 000000000000..4ad37f78f4ab --- /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 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 main = new TupleTag() {}; + TupleTag other = new TupleTag() {}; + assertUnsupported( + kv("o") + .apply(ParDo.of(new StatelessDoFn()).withOutputTags(main, TupleTagList.of(other))) + .getPipeline(), + "additional outputs"); + } + + private static final class ToKvFn extends DoFn> { + @ProcessElement + public void process(@Element String element, OutputReceiver> out) { + out.output(KV.of(element, 1)); + } + } + + private static final class StatelessDoFn extends DoFn, Integer> { + @ProcessElement + public void process() {} + } + + private static final class StatefulDoFn extends DoFn, 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 000000000000..86369226ff3b --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatelessParDoStreamingTest.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 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.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; + +/** 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 { + @ProcessElement + public void process(@Element String element, OutputReceiver out) { + out.output(element); + } + } + + @After + public void tearDown() { + TestUnboundedSource.forget("stateless-pardo"); + TestUnboundedSource.forget("flatten-a"); + TestUnboundedSource.forget("flatten-b"); + } + + @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 collected = new HashSet<>(StreamingTestUtils.getCollected(collectorId)); + Set 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 a = pipeline.apply("ReadA", Read.from(new TestUnboundedSource(tagA, 1, 5))); + PCollection 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 collected = new HashSet<>(StreamingTestUtils.getCollected(collectorId)); + Set 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 000000000000..e855050b1e69 --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingCheckpointRestartTest.java @@ -0,0 +1,120 @@ +/* + * 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.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +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. + * + *

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); + + List collectedA = new ArrayList<>(StreamingTestUtils.getCollected(collectorA)); + assertEquals( + "first run must read every element", + TestUnboundedSource.elements(TAG, 1, ELEMENT_COUNT), + new HashSet<>(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); + + List collectedB = new ArrayList<>(StreamingTestUtils.getCollected(collectorB)); + assertTrue( + "second run must not re-emit elements the first run committed", + Collections.disjoint(new HashSet<>(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 000000000000..5592cec594ab --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java @@ -0,0 +1,193 @@ +/* + * 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.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.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; + +/** + * State transitions for a streaming pipeline. + * + *

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(); + + /** How long to wait for a query to start before failing. */ + private static final long QUERY_START_TIMEOUT_MILLIS = 60_000L; + + @After + public void tearDown() { + TestUnboundedSource.forget("lifecycle-done"); + TestUnboundedSource.forget("lifecycle-cancel"); + TestUnboundedSource.forget("lifecycle-healthy"); + TestUnboundedSource.forget("lifecycle-poison"); + } + + /** Blocks until at least one streaming query is active on the shared session. */ + private static void awaitQueryStarted() throws InterruptedException { + long deadline = System.currentTimeMillis() + QUERY_START_TIMEOUT_MILLIS; + while (SESSION.getSession().streams().active().length == 0) { + assertTrue( + "no streaming query started 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 collected = new HashSet<>(StreamingTestUtils.getCollected(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()); + + awaitQueryStarted(); + + PipelineResult.State cancelledState = result.cancel(); + assertEquals(PipelineResult.State.CANCELLED, cancelledState); + assertEquals(PipelineResult.State.CANCELLED, result.getState()); + + long deadline = System.currentTimeMillis() + QUERY_START_TIMEOUT_MILLIS; + while (SESSION.getSession().streams().active().length > 0) { + assertTrue( + "the streaming query was still active " + QUERY_START_TIMEOUT_MILLIS + "ms after cancel", + System.currentTimeMillis() < deadline); + Thread.sleep(50L); + } + } + + /** 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()); + + long deadline = System.currentTimeMillis() + QUERY_START_TIMEOUT_MILLIS; + while (SESSION.getSession().streams().active().length > 0) { + assertTrue( + "a sibling query was still active " + + QUERY_START_TIMEOUT_MILLIS + + "ms after the pipeline failed", + System.currentTimeMillis() < deadline); + Thread.sleep(50L); + } + } + + /** Throws on one specific element index, passes every other element through. */ + private static final class ThrowOnElementDoFn extends DoFn { + private final int poisonIndex; + + ThrowOnElementDoFn(int poisonIndex) { + this.poisonIndex = poisonIndex; + } + + @ProcessElement + public void processElement(@Element String element, OutputReceiver 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 000000000000..cd5bf911575e --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTestUtils.java @@ -0,0 +1,154 @@ +/* + * 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.List; +import java.util.Map; +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. + * + *

Collectors are static so they work in local mode only. + * + *

The {@link #run} helper bounds each query at five minutes and cancels on expiry. + * + *

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> 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 extends DoFn { + private final String collectorId; + + public CollectDoFn(String collectorId) { + this.collectorId = Preconditions.checkNotNull(collectorId); + } + + @ProcessElement + public void processElement(@Element T element, OutputReceiver 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 List getCollected(String collectorId) { + List values = COLLECTORS.get(collectorId); + if (values == null) { + return Collections.emptyList(); + } + synchronized (values) { + return (List) new ArrayList<>(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; + } +} From 75dbe36e66970553347386f723b3525a9d5a6110 Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Fri, 11 Sep 2026 07:35:13 +0000 Subject: [PATCH 2/6] [Spark] Let cancel() wait for the execution thread before stopping the session SparkStructuredStreamingPipelineResult.cancel() interrupted the execution thread and ran the terminal state callback at once, which stops the SparkSession. The thread kept translating or evaluating on a stopped SparkContext. When that happened during the static initialization of PipelineTranslatorBatch the class was poisoned for the JVM and every later batch pipeline failed with NoClassDefFoundError, seen in the Spark Versions PreCommit on #40090. runAsync now hands the result its single thread executor and cancel() waits for it to terminate, bounded at 60 seconds, before the callback runs. --- ...parkStructuredStreamingPipelineResult.java | 29 +++++- .../SparkStructuredStreamingRunner.java | 22 ++++- ...StructuredStreamingPipelineResultTest.java | 91 +++++++++++++++++++ 3 files changed, 136 insertions(+), 6 deletions(-) create mode 100644 runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResultTest.java diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResult.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResult.java index b592b6fb742d..14fadd55fc4f 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResult.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResult.java @@ -22,6 +22,7 @@ import java.io.IOException; import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; @@ -35,25 +36,36 @@ import org.apache.spark.SparkException; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Duration; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; public class SparkStructuredStreamingPipelineResult implements PipelineResult { + private static final Logger LOG = + LoggerFactory.getLogger(SparkStructuredStreamingPipelineResult.class); + + /** Upper bound on how long {@link #cancel()} waits for the execution thread to end. */ + private static final long CANCEL_WAIT_SECONDS = 60; + private final Future pipelineExecution; // Supplies the context of the translated pipeline, null until translation has completed. private final Supplier evaluationContext; private final MetricsAccumulator metrics; private final @Nullable Runnable onTerminalState; + private final ExecutorService executor; private PipelineResult.State state; SparkStructuredStreamingPipelineResult( Future pipelineExecution, Supplier evaluationContext, MetricsAccumulator metrics, - final @Nullable Runnable onTerminalState) { + final @Nullable Runnable onTerminalState, + ExecutorService executor) { this.pipelineExecution = pipelineExecution; this.evaluationContext = evaluationContext; this.metrics = metrics; this.onTerminalState = onTerminalState; + this.executor = executor; // pipelineExecution is expected to have started executing eagerly. this.state = State.RUNNING; } @@ -117,6 +129,11 @@ public MetricResults metrics() { return asAttemptedOnlyMetricResults(metrics.value()); } + /** + * Cancels the execution and waits up to {@link #CANCEL_WAIT_SECONDS} for the execution thread to + * end before the terminal state callback stops the session. An interrupted caller returns without + * the callback, the state stays RUNNING. + */ @Override public PipelineResult.State cancel() throws IOException { EvaluationContext ctx = evaluationContext.get(); @@ -124,6 +141,16 @@ public PipelineResult.State cancel() throws IOException { ctx.stop(); } pipelineExecution.cancel(true); + try { + if (!executor.awaitTermination(CANCEL_WAIT_SECONDS, TimeUnit.SECONDS)) { + LOG.warn( + "Pipeline execution still running {} s after cancel, stopping the session anyway.", + CANCEL_WAIT_SECONDS); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return state; + } offerNewState(PipelineResult.State.CANCELLED); return state; } diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingRunner.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingRunner.java index f78026847fad..8cce4c18e866 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingRunner.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingRunner.java @@ -152,7 +152,7 @@ public SparkStructuredStreamingPipelineResult run(final Pipeline pipeline) { // evaluation on cancel. Remains null until translation completes. final AtomicReference ctxRef = new AtomicReference<>(); - final Future submissionFuture = + final Submission submission = runAsync( () -> { EvaluationContext ctx = translatePipeline(sparkSession, pipeline); @@ -162,10 +162,11 @@ public SparkStructuredStreamingPipelineResult run(final Pipeline pipeline) { final SparkStructuredStreamingPipelineResult result = new SparkStructuredStreamingPipelineResult( - submissionFuture, + submission.future, ctxRef::get, metrics, - sparkStopFn(sparkSession, options.getUseActiveSparkSession())); + sparkStopFn(sparkSession, options.getUseActiveSparkSession()), + submission.executor); if (options.getEnableSparkMetricSinks()) { registerMetricsSource(options.getAppName(), metrics); @@ -217,7 +218,18 @@ private void startMetricsPusher( } } - private static Future runAsync(Runnable task) { + /** The future of a submitted pipeline and the executor that runs it, for cancel to await. */ + private static final class Submission { + private final Future future; + private final ExecutorService executor; + + Submission(Future future, ExecutorService executor) { + this.future = future; + this.executor = executor; + } + } + + private static Submission runAsync(Runnable task) { ThreadFactory factory = new ThreadFactoryBuilder() .setDaemon(true) @@ -226,7 +238,7 @@ private static Future runAsync(Runnable task) { ExecutorService execService = Executors.newSingleThreadExecutor(factory); Future future = execService.submit(task); execService.shutdown(); - return future; + return new Submission(future, execService); } private static @Nullable Runnable sparkStopFn(SparkSession session, boolean isProvided) { diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResultTest.java b/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResultTest.java new file mode 100644 index 000000000000..0ef31df3c4a9 --- /dev/null +++ b/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResultTest.java @@ -0,0 +1,91 @@ +/* + * 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; + +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.Uninterruptibles.sleepUninterruptibly; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import org.apache.beam.runners.spark.structuredstreaming.metrics.MetricsAccumulator; +import org.apache.beam.sdk.PipelineResult; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for {@link SparkStructuredStreamingPipelineResult#cancel()}. */ +@RunWith(JUnit4.class) +public class SparkStructuredStreamingPipelineResultTest { + + private static SparkStructuredStreamingPipelineResult result( + Future f, Runnable cb, ExecutorService exec) { + return new SparkStructuredStreamingPipelineResult( + f, () -> null, new MetricsAccumulator(), cb, exec); + } + + private static Future submit(ExecutorService exec, Runnable r) throws Exception { + CountDownLatch started = new CountDownLatch(1); + Future f = + exec.submit( + () -> { + started.countDown(); + r.run(); + }); + exec.shutdown(); + started.await(); + return f; + } + + @Test + public void testCancelWaitsForExecutionThread() throws Exception { + CountDownLatch block = new CountDownLatch(1); + AtomicBoolean finished = new AtomicBoolean(); + AtomicBoolean observed = new AtomicBoolean(); + ExecutorService exec = Executors.newSingleThreadExecutor(); + Future f = + submit( + exec, + () -> { + try { + block.await(); + } catch (InterruptedException ignored) { + } + sleepUninterruptibly(300, TimeUnit.MILLISECONDS); + finished.set(true); + }); + SparkStructuredStreamingPipelineResult res = + result(f, () -> observed.set(finished.get()), exec); + assertEquals(PipelineResult.State.CANCELLED, res.cancel()); + assertTrue(observed.get()); + assertEquals(PipelineResult.State.CANCELLED, res.getState()); + } + + @Test + public void testCancelAwaitsTaskIgnoringInterrupt() throws Exception { + ExecutorService exec = Executors.newSingleThreadExecutor(); + Future f = submit(exec, () -> sleepUninterruptibly(200, TimeUnit.MILLISECONDS)); + SparkStructuredStreamingPipelineResult res = result(f, null, exec); + assertEquals(PipelineResult.State.CANCELLED, res.cancel()); + assertTrue(exec.isTerminated()); + } +} From 58303d6da346da4c22f8351a724e011b5fa636bf Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Fri, 11 Sep 2026 08:04:55 +0000 Subject: [PATCH 3/6] [Spark 4] Address review of the stateless streaming translator Split the translator registry into PipelineTranslatorCommon with a thin PipelineTranslatorBatch subclass, the streaming translator extends the common class. Exceptions are created at the call site, the unchecked cast is scoped to one statement, the checkpoint location is joined with a Path, a failing evaluate stops every started query, and stop failures log the exception. Tests share one polling helper, collect results as one Set snapshot, and only the restart test keeps a source counter teardown. --- .../PipelineTranslatorStreaming.java | 48 ++++---- .../StreamingEvaluationContext.java | 13 ++- .../StatelessParDoStreamingTest.java | 12 +- .../StreamingCheckpointRestartTest.java | 12 +- .../StreamingPipelineLifecycleTest.java | 44 ++----- .../streaming/StreamingTestUtils.java | 8 +- .../batch/ParDoTranslatorBatch.java | 2 +- .../batch/PipelineTranslatorBatch.java | 91 +-------------- .../batch/PipelineTranslatorCommon.java | 110 ++++++++++++++++++ .../batch/StatefulParDoTranslatorBatch.java | 2 +- 10 files changed, 173 insertions(+), 169 deletions(-) create mode 100644 runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/PipelineTranslatorCommon.java 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 d0ab3c25b615..e76d0f1083b7 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 @@ -19,7 +19,7 @@ import java.util.Collection; import org.apache.beam.runners.spark.SparkCommonPipelineOptions; -import org.apache.beam.runners.spark.structuredstreaming.translation.batch.PipelineTranslatorBatch; +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; @@ -36,66 +36,74 @@ import org.checkerframework.checker.nullness.qual.Nullable; /** - * Pipeline translator for streaming pipelines on Spark 4. It extends the batch translator to reuse + * 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 PipelineTranslatorBatch { +public class PipelineTranslatorStreaming extends PipelineTranslatorCommon { + + 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 - @SuppressWarnings({"rawtypes", "unchecked"}) @Nullable protected > TransformTranslator getTransformTranslator(TransformT transform) { if (transform instanceof SplittableParDo.PrimitiveUnboundedRead) { - return (TransformTranslator) new ReadUnboundedTranslator<>(); + @SuppressWarnings("unchecked") + TransformTranslator read = + (TransformTranslator) + (TransformTranslator) new ReadUnboundedTranslator<>(); + return read; } if (transform instanceof SplittableParDo.PrimitiveBoundedRead) { - throw unsupported( - "Bounded Read (Read.from(BoundedSource), Create with two or more elements)"); + throw new UnsupportedOperationException( + "Bounded Read (Read.from(BoundedSource), Create with two or more elements)" + + NOT_SUPPORTED); } if (transform instanceof Impulse) { - throw unsupported("Impulse (Create with fewer than two elements, PAssert)"); + throw new UnsupportedOperationException( + "Impulse (Create with fewer than two elements, PAssert)" + NOT_SUPPORTED); } if (transform instanceof GroupByKey) { - throw unsupported("GroupByKey"); + throw new UnsupportedOperationException("GroupByKey" + NOT_SUPPORTED); } if (transform instanceof Combine.PerKey) { - throw unsupported("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 unsupported("Stateful ParDo (" + signature.fnClass().getName() + ")"); + throw new UnsupportedOperationException( + "Stateful ParDo (" + signature.fnClass().getName() + ")" + NOT_SUPPORTED); } if (!parDo.getSideInputs().isEmpty()) { - throw unsupported("ParDo with side inputs (" + signature.fnClass().getName() + ")"); + throw new UnsupportedOperationException( + "ParDo with side inputs (" + signature.fnClass().getName() + ")" + NOT_SUPPORTED); } if (!parDo.getAdditionalOutputTags().getAll().isEmpty()) { - throw unsupported("ParDo with additional outputs (" + signature.fnClass().getName() + ")"); + throw new UnsupportedOperationException( + "ParDo with additional outputs (" + + signature.fnClass().getName() + + ")" + + NOT_SUPPORTED); } } return super.getTransformTranslator(transform); } - private static UnsupportedOperationException unsupported(String what) { - return new UnsupportedOperationException( - what - + " is not supported by the Spark 4 streaming runner yet, see" - + " https://github.com/apache/beam/issues/36841"); - } - @Override protected EvaluationContext createEvaluationContext( Collection> leaves, 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 index 1edf3ea748cd..ff7839bddb20 100644 --- 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 @@ -31,6 +31,7 @@ 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; @@ -116,6 +117,9 @@ public void evaluate() { toAwait = new ArrayList<>(queries); } awaitTermination(toAwait); + } catch (RuntimeException e) { + stop(); + throw e; } finally { if (idleStopListener != null) { getSparkSession().streams().removeListener(idleStopListener); @@ -151,7 +155,9 @@ private StreamingQuery startQuery( .writeStream() .format("noop") .outputMode("append") - .option("checkpointLocation", checkpointBaseDir + "/" + leafIndex) + .option( + "checkpointLocation", + new Path(checkpointBaseDir, Integer.toString(leafIndex)).toString()) .trigger(Trigger.ProcessingTime(options.getMaxBatchDurationMillis())) .start(); } catch (TimeoutException e) { @@ -187,10 +193,7 @@ private void stopQuery(StreamingQuery query) { query.stop(); } } catch (TimeoutException | RuntimeException e) { - LOG.warn( - "Error while stopping streaming query {}: {}", - query.id(), - String.valueOf(e.getMessage())); + LOG.warn("Failed to stop streaming query {}.", query.id(), e); } } 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 index 86369226ff3b..d41f5e774227 100644 --- 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 @@ -34,7 +34,6 @@ import org.apache.beam.sdk.transforms.ParDo; import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.PCollectionList; -import org.junit.After; import org.junit.ClassRule; import org.junit.Rule; import org.junit.Test; @@ -59,13 +58,6 @@ public void process(@Element String element, OutputReceiver out) { } } - @After - public void tearDown() { - TestUnboundedSource.forget("stateless-pardo"); - TestUnboundedSource.forget("flatten-a"); - TestUnboundedSource.forget("flatten-b"); - } - @Test public void everyElementPassesThrough() throws Exception { String tag = "stateless-pardo"; @@ -83,7 +75,7 @@ public void everyElementPassesThrough() throws Exception { PipelineResult result = StreamingTestUtils.run(pipeline); - Set collected = new HashSet<>(StreamingTestUtils.getCollected(collectorId)); + Set collected = StreamingTestUtils.collected(collectorId); Set expected = TestUnboundedSource.elements(tag, 1, 10); assertEquals("pipeline state=" + result.getState(), expected, collected); } @@ -109,7 +101,7 @@ public void flattenCombinesMultipleUnboundedSources() throws Exception { PipelineResult result = StreamingTestUtils.run(pipeline); - Set collected = new HashSet<>(StreamingTestUtils.getCollected(collectorId)); + Set collected = StreamingTestUtils.collected(collectorId); Set expected = new HashSet<>(); expected.addAll(TestUnboundedSource.elements(tagA, 1, 5)); expected.addAll(TestUnboundedSource.elements(tagB, 1, 5)); 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 index e855050b1e69..c7ad2fa62956 100644 --- 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 @@ -22,10 +22,8 @@ import java.io.File; import java.io.Serializable; -import java.util.ArrayList; import java.util.Collections; -import java.util.HashSet; -import java.util.List; +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; @@ -78,11 +76,11 @@ public void restartedPipelineResumesFromDurableCheckpointMarks() throws Exceptio // First run reads all elements from a fresh checkpoint directory. runPipeline(checkpointPath, collectorA, TAG); - List collectedA = new ArrayList<>(StreamingTestUtils.getCollected(collectorA)); + Set collectedA = StreamingTestUtils.collected(collectorA); assertEquals( "first run must read every element", TestUnboundedSource.elements(TAG, 1, ELEMENT_COUNT), - new HashSet<>(collectedA)); + collectedA); // The source checkpoint lives under the location the translator handed to Spark. File sourceRoot = new File(new File(checkpointPath, "0"), "sources/0"); @@ -96,10 +94,10 @@ public void restartedPipelineResumesFromDurableCheckpointMarks() throws Exceptio // Second run resumes against the same checkpoint directory. runPipeline(checkpointPath, collectorB, TAG); - List collectedB = new ArrayList<>(StreamingTestUtils.getCollected(collectorB)); + Set collectedB = StreamingTestUtils.collected(collectorB); assertTrue( "second run must not re-emit elements the first run committed", - Collections.disjoint(new HashSet<>(collectedA), collectedB)); + Collections.disjoint(collectedA, collectedB)); assertTrue( "readers must be recreated during the second run", TestUnboundedSource.created(TAG) > createdBeforeSecondRun); 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 index 5592cec594ab..76462c38fe10 100644 --- 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 @@ -22,8 +22,8 @@ import static org.junit.Assert.assertTrue; import java.io.Serializable; -import java.util.HashSet; 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; @@ -33,7 +33,6 @@ import org.apache.beam.sdk.io.Read; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.ParDo; -import org.junit.After; import org.junit.ClassRule; import org.junit.Rule; import org.junit.Test; @@ -56,23 +55,16 @@ public class StreamingPipelineLifecycleTest implements Serializable { @Rule public transient TemporaryFolder checkpointDir = new TemporaryFolder(); - /** How long to wait for a query to start before failing. */ + /** Upper bound for a query to start or stop. */ private static final long QUERY_START_TIMEOUT_MILLIS = 60_000L; - @After - public void tearDown() { - TestUnboundedSource.forget("lifecycle-done"); - TestUnboundedSource.forget("lifecycle-cancel"); - TestUnboundedSource.forget("lifecycle-healthy"); - TestUnboundedSource.forget("lifecycle-poison"); - } - - /** Blocks until at least one streaming query is active on the shared session. */ - private static void awaitQueryStarted() throws InterruptedException { + /** 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 (SESSION.getSession().streams().active().length == 0) { + while (!condition.test(SESSION.getSession().streams().active().length)) { assertTrue( - "no streaming query started within " + QUERY_START_TIMEOUT_MILLIS + "ms", + failure + " within " + QUERY_START_TIMEOUT_MILLIS + " ms", System.currentTimeMillis() < deadline); Thread.sleep(50L); } @@ -99,7 +91,7 @@ public void idlePipelineGoesFromRunningToDoneOnceIdle() throws Exception { assertEquals(PipelineResult.State.DONE, finalState); assertEquals(PipelineResult.State.DONE, result.getState()); - Set collected = new HashSet<>(StreamingTestUtils.getCollected(collectorId)); + Set collected = StreamingTestUtils.collected(collectorId); assertEquals(TestUnboundedSource.elements(tag, 1, 10), collected); } @@ -122,19 +114,13 @@ public void cancelStopsTheQueryAndReportsCancelled() throws Exception { PipelineResult result = pipeline.run(); assertEquals(PipelineResult.State.RUNNING, result.getState()); - awaitQueryStarted(); + awaitActiveQueries(count -> count > 0, "no streaming query started"); PipelineResult.State cancelledState = result.cancel(); assertEquals(PipelineResult.State.CANCELLED, cancelledState); assertEquals(PipelineResult.State.CANCELLED, result.getState()); - long deadline = System.currentTimeMillis() + QUERY_START_TIMEOUT_MILLIS; - while (SESSION.getSession().streams().active().length > 0) { - assertTrue( - "the streaming query was still active " + QUERY_START_TIMEOUT_MILLIS + "ms after cancel", - System.currentTimeMillis() < deadline); - Thread.sleep(50L); - } + awaitActiveQueries(count -> count == 0, "the streaming query did not stop after cancel"); } /** A failure in any leaf query surfaces through waitUntilFinish. */ @@ -163,15 +149,7 @@ public void failingLeafQueryFailsThePipelineAndStopsHealthySibling() throws Exce assertThrows(RuntimeException.class, () -> StreamingTestUtils.waitUntilFinish(result)); assertEquals(PipelineResult.State.FAILED, result.getState()); - long deadline = System.currentTimeMillis() + QUERY_START_TIMEOUT_MILLIS; - while (SESSION.getSession().streams().active().length > 0) { - assertTrue( - "a sibling query was still active " - + QUERY_START_TIMEOUT_MILLIS - + "ms after the pipeline failed", - System.currentTimeMillis() < deadline); - Thread.sleep(50L); - } + awaitActiveQueries(count -> count == 0, "a sibling query did not stop after the failure"); } /** Throws on one specific element index, passes every other element through. */ 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 cd5bf911575e..82ecd7dd19f6 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 @@ -20,8 +20,10 @@ 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; @@ -77,13 +79,13 @@ private static void append(String collectorId, Object value) { /** Returns a snapshot of everything collected so far under {@code collectorId}. */ @SuppressWarnings("unchecked") - public static List getCollected(String collectorId) { + public static Set collected(String collectorId) { List values = COLLECTORS.get(collectorId); if (values == null) { - return Collections.emptyList(); + return Collections.emptySet(); } synchronized (values) { - return (List) new ArrayList<>(values); + return (Set) (Set) new HashSet<>(values); } } 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 14058a73733d..b454fc1a34ef 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 @@ public boolean canTranslate(ParDo.MultiOutput transform) { 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 ba7cbb0fa035..df8977063cd6 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, 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. - * - *

Currently only {@link ParDo.MultiOutput} needs this, to route stateful and time sorted - * {@link org.apache.beam.sdk.transforms.DoFn DoFns} away from {@link ParDoTranslatorBatch}. - */ - @SuppressWarnings("rawtypes") - private static final TransformTranslator STATEFUL_PARDO_TRANSLATOR = - new StatefulParDoTranslatorBatch<>(); - - /** Returns a {@link TransformTranslator} for the given {@link PTransform} if known. */ - @Override - @Nullable - protected > - TransformTranslator getTransformTranslator(TransformT transform) { - // 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/PipelineTranslatorCommon.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/PipelineTranslatorCommon.java new file mode 100644 index 000000000000..b71d35967605 --- /dev/null +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/PipelineTranslatorCommon.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.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; + +/** + * Holds the translator registry shared by the batch and the streaming translator, and the + * registry lookup. + */ +@Internal +public class PipelineTranslatorCommon extends PipelineTranslator { + + // -------------------------------------------------------------------------------------------- + // Transform Translator Registry + // -------------------------------------------------------------------------------------------- + + @SuppressWarnings("rawtypes") + private static final Map, 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. + * + *

Currently only {@link ParDo.MultiOutput} needs this, to route stateful and time sorted + * {@link org.apache.beam.sdk.transforms.DoFn DoFns} away from {@link ParDoTranslatorBatch}. + */ + @SuppressWarnings("rawtypes") + private static final TransformTranslator STATEFUL_PARDO_TRANSLATOR = + new StatefulParDoTranslatorBatch<>(); + + /** Returns a {@link TransformTranslator} for the given {@link PTransform} if known. */ + @Override + @Nullable + protected > + TransformTranslator getTransformTranslator(TransformT transform) { + // Resolved ahead of the class keyed registry: ParDo.MultiOutput maps to a different translator + // depending on the DoFn signature, which a lookup by transform class alone cannot express. This + // is the predicated dispatch of the TODO above, limited to the single transform needing it. + if (transform instanceof ParDo.MultiOutput + && StatefulParDoTranslatorBatch.appliesTo((ParDo.MultiOutput) transform)) { + return STATEFUL_PARDO_TRANSLATOR; + } + return TRANSFORM_TRANSLATORS.get(transform.getClass()); + } +} diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/StatefulParDoTranslatorBatch.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/StatefulParDoTranslatorBatch.java index 75d84b630bb8..5f060b6aef41 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 @@ /** * Translator for a stateful {@link ParDo.MultiOutput}, or one requiring time sorted input. * - *

Selected by {@link PipelineTranslatorBatch} in place of {@link ParDoTranslatorBatch} when the + *

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}. * From 9217c1ad5c4653ee8e2b80d2953ab0045ad4cd1d Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Fri, 11 Sep 2026 08:17:10 +0000 Subject: [PATCH 4/6] [Spark] Apply spotless to the shared translator module --- .../translation/batch/PipelineTranslatorCommon.java | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/PipelineTranslatorCommon.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/PipelineTranslatorCommon.java index b71d35967605..899e0443a94f 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/PipelineTranslatorCommon.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/PipelineTranslatorCommon.java @@ -21,7 +21,6 @@ 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,8 +36,8 @@ import org.checkerframework.checker.nullness.qual.Nullable; /** - * Holds the translator registry shared by the batch and the streaming translator, and the - * registry lookup. + * Holds the translator registry shared by the batch and the streaming translator, and the registry + * lookup. */ @Internal public class PipelineTranslatorCommon extends PipelineTranslator { From 8bad262324ca663341eac0a15233d1676ba44cf5 Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Fri, 11 Sep 2026 16:26:52 +0000 Subject: [PATCH 5/6] Revert "[Spark] Let cancel() wait for the execution thread before stopping the session" This reverts commit 75dbe36e66970553347386f723b3525a9d5a6110. --- ...parkStructuredStreamingPipelineResult.java | 29 +----- .../SparkStructuredStreamingRunner.java | 22 +---- ...StructuredStreamingPipelineResultTest.java | 91 ------------------- 3 files changed, 6 insertions(+), 136 deletions(-) delete mode 100644 runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResultTest.java diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResult.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResult.java index 14fadd55fc4f..b592b6fb742d 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResult.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResult.java @@ -22,7 +22,6 @@ import java.io.IOException; import java.util.concurrent.ExecutionException; -import java.util.concurrent.ExecutorService; import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; @@ -36,36 +35,25 @@ import org.apache.spark.SparkException; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Duration; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; public class SparkStructuredStreamingPipelineResult implements PipelineResult { - private static final Logger LOG = - LoggerFactory.getLogger(SparkStructuredStreamingPipelineResult.class); - - /** Upper bound on how long {@link #cancel()} waits for the execution thread to end. */ - private static final long CANCEL_WAIT_SECONDS = 60; - private final Future pipelineExecution; // Supplies the context of the translated pipeline, null until translation has completed. private final Supplier evaluationContext; private final MetricsAccumulator metrics; private final @Nullable Runnable onTerminalState; - private final ExecutorService executor; private PipelineResult.State state; SparkStructuredStreamingPipelineResult( Future pipelineExecution, Supplier evaluationContext, MetricsAccumulator metrics, - final @Nullable Runnable onTerminalState, - ExecutorService executor) { + final @Nullable Runnable onTerminalState) { this.pipelineExecution = pipelineExecution; this.evaluationContext = evaluationContext; this.metrics = metrics; this.onTerminalState = onTerminalState; - this.executor = executor; // pipelineExecution is expected to have started executing eagerly. this.state = State.RUNNING; } @@ -129,11 +117,6 @@ public MetricResults metrics() { return asAttemptedOnlyMetricResults(metrics.value()); } - /** - * Cancels the execution and waits up to {@link #CANCEL_WAIT_SECONDS} for the execution thread to - * end before the terminal state callback stops the session. An interrupted caller returns without - * the callback, the state stays RUNNING. - */ @Override public PipelineResult.State cancel() throws IOException { EvaluationContext ctx = evaluationContext.get(); @@ -141,16 +124,6 @@ public PipelineResult.State cancel() throws IOException { ctx.stop(); } pipelineExecution.cancel(true); - try { - if (!executor.awaitTermination(CANCEL_WAIT_SECONDS, TimeUnit.SECONDS)) { - LOG.warn( - "Pipeline execution still running {} s after cancel, stopping the session anyway.", - CANCEL_WAIT_SECONDS); - } - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - return state; - } offerNewState(PipelineResult.State.CANCELLED); return state; } diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingRunner.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingRunner.java index 8cce4c18e866..f78026847fad 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingRunner.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingRunner.java @@ -152,7 +152,7 @@ public SparkStructuredStreamingPipelineResult run(final Pipeline pipeline) { // evaluation on cancel. Remains null until translation completes. final AtomicReference ctxRef = new AtomicReference<>(); - final Submission submission = + final Future submissionFuture = runAsync( () -> { EvaluationContext ctx = translatePipeline(sparkSession, pipeline); @@ -162,11 +162,10 @@ public SparkStructuredStreamingPipelineResult run(final Pipeline pipeline) { final SparkStructuredStreamingPipelineResult result = new SparkStructuredStreamingPipelineResult( - submission.future, + submissionFuture, ctxRef::get, metrics, - sparkStopFn(sparkSession, options.getUseActiveSparkSession()), - submission.executor); + sparkStopFn(sparkSession, options.getUseActiveSparkSession())); if (options.getEnableSparkMetricSinks()) { registerMetricsSource(options.getAppName(), metrics); @@ -218,18 +217,7 @@ private void startMetricsPusher( } } - /** The future of a submitted pipeline and the executor that runs it, for cancel to await. */ - private static final class Submission { - private final Future future; - private final ExecutorService executor; - - Submission(Future future, ExecutorService executor) { - this.future = future; - this.executor = executor; - } - } - - private static Submission runAsync(Runnable task) { + private static Future runAsync(Runnable task) { ThreadFactory factory = new ThreadFactoryBuilder() .setDaemon(true) @@ -238,7 +226,7 @@ private static Submission runAsync(Runnable task) { ExecutorService execService = Executors.newSingleThreadExecutor(factory); Future future = execService.submit(task); execService.shutdown(); - return new Submission(future, execService); + return future; } private static @Nullable Runnable sparkStopFn(SparkSession session, boolean isProvided) { diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResultTest.java b/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResultTest.java deleted file mode 100644 index 0ef31df3c4a9..000000000000 --- a/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResultTest.java +++ /dev/null @@ -1,91 +0,0 @@ -/* - * 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; - -import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.Uninterruptibles.sleepUninterruptibly; -import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertTrue; - -import java.util.concurrent.CountDownLatch; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; -import java.util.concurrent.Future; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicBoolean; -import org.apache.beam.runners.spark.structuredstreaming.metrics.MetricsAccumulator; -import org.apache.beam.sdk.PipelineResult; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; - -/** Tests for {@link SparkStructuredStreamingPipelineResult#cancel()}. */ -@RunWith(JUnit4.class) -public class SparkStructuredStreamingPipelineResultTest { - - private static SparkStructuredStreamingPipelineResult result( - Future f, Runnable cb, ExecutorService exec) { - return new SparkStructuredStreamingPipelineResult( - f, () -> null, new MetricsAccumulator(), cb, exec); - } - - private static Future submit(ExecutorService exec, Runnable r) throws Exception { - CountDownLatch started = new CountDownLatch(1); - Future f = - exec.submit( - () -> { - started.countDown(); - r.run(); - }); - exec.shutdown(); - started.await(); - return f; - } - - @Test - public void testCancelWaitsForExecutionThread() throws Exception { - CountDownLatch block = new CountDownLatch(1); - AtomicBoolean finished = new AtomicBoolean(); - AtomicBoolean observed = new AtomicBoolean(); - ExecutorService exec = Executors.newSingleThreadExecutor(); - Future f = - submit( - exec, - () -> { - try { - block.await(); - } catch (InterruptedException ignored) { - } - sleepUninterruptibly(300, TimeUnit.MILLISECONDS); - finished.set(true); - }); - SparkStructuredStreamingPipelineResult res = - result(f, () -> observed.set(finished.get()), exec); - assertEquals(PipelineResult.State.CANCELLED, res.cancel()); - assertTrue(observed.get()); - assertEquals(PipelineResult.State.CANCELLED, res.getState()); - } - - @Test - public void testCancelAwaitsTaskIgnoringInterrupt() throws Exception { - ExecutorService exec = Executors.newSingleThreadExecutor(); - Future f = submit(exec, () -> sleepUninterruptibly(200, TimeUnit.MILLISECONDS)); - SparkStructuredStreamingPipelineResult res = result(f, null, exec); - assertEquals(PipelineResult.State.CANCELLED, res.cancel()); - assertTrue(exec.isTerminated()); - } -} From 8b5927d0e977250de59458ef2cc08c7b227473ba Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Fri, 11 Sep 2026 16:33:31 +0000 Subject: [PATCH 6/6] [Spark 4] Drop the redundant casts in the streaming translator and test utils --- .../translation/PipelineTranslatorStreaming.java | 9 ++++----- .../translation/streaming/StreamingTestUtils.java | 2 +- 2 files changed, 5 insertions(+), 6 deletions(-) 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 e76d0f1083b7..3529daaff313 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 @@ -44,6 +44,9 @@ @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"; @@ -55,11 +58,7 @@ public class PipelineTranslatorStreaming extends PipelineTranslatorCommon { TransformTranslator getTransformTranslator(TransformT transform) { if (transform instanceof SplittableParDo.PrimitiveUnboundedRead) { - @SuppressWarnings("unchecked") - TransformTranslator read = - (TransformTranslator) - (TransformTranslator) new ReadUnboundedTranslator<>(); - return read; + return READ_UNBOUNDED; } if (transform instanceof SplittableParDo.PrimitiveBoundedRead) { 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 82ecd7dd19f6..3d7841154150 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 @@ -85,7 +85,7 @@ public static Set collected(String collectorId) { return Collections.emptySet(); } synchronized (values) { - return (Set) (Set) new HashSet<>(values); + return (Set) new HashSet<>(values); } }