tkaymak commented on code in PR #40090: URL: https://github.com/apache/beam/pull/40090#discussion_r3987110093
########## 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 { Review Comment: Done. The registry and the lookup live in PipelineTranslatorCommon now, PipelineTranslatorBatch is a thin subclass and the streaming translator extends the common class. ########## 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 <InT extends PInput, OutT extends POutput, TransformT extends PTransform<InT, OutT>> + TransformTranslator<InT, OutT, TransformT> 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) { Review Comment: Done, the helper is gone, each site throws its own UnsupportedOperationException. ########## 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"}) Review Comment: Done. One statement level unchecked cast for the read translator remains, no raw types. ########## 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. + * + * <p>Leaf {@code i} checkpoints under {@code <checkpointDir>/i}, in pipeline graph order. A changed + * pipeline needs a new checkpoint directory, as with any Spark streaming query. + */ +@Internal +public class StreamingEvaluationContext extends EvaluationContext { + private static final Logger LOG = LoggerFactory.getLogger(StreamingEvaluationContext.class); + + private static final long AWAIT_POLL_TIMEOUT_MILLIS = 100; + + private final SparkStructuredStreamingPipelineOptions options; + + // Guards queries and stopped. + private final Object lock = new Object(); + private final List<StreamingQuery> queries = new ArrayList<>(); + private boolean stopped = false; + + StreamingEvaluationContext( + Collection<? extends NamedDataset<?>> leaves, + SparkSession session, + SparkCommonPipelineOptions options) { + super(leaves, session); + this.options = options.as(SparkStructuredStreamingPipelineOptions.class); + } + + /** Starts one streaming query per leaf dataset and blocks until all queries terminate. */ + @Override + public void evaluate() { + String checkpointBaseDir = options.getCheckpointDir(); + checkArgument( + checkpointBaseDir != null && !checkpointBaseDir.isEmpty(), + "checkpointDir must be set for a streaming pipeline"); + int idleStopThreshold = options.getStreamingStopAfterIdleBatches(); + + StreamingQueryListener idleStopListener = null; + if (idleStopThreshold >= 0) { + idleStopListener = new IdleStopListener(idleStopThreshold); + getSparkSession().streams().addListener(idleStopListener); + } + + try { + int leafIndex = 0; + for (NamedDataset<?> ds : leaves()) { + Dataset<?> dataset = ds.dataset(); + if (dataset == null) { + continue; + } + synchronized (lock) { + if (stopped) { + break; + } + } + if (!dataset.isStreaming()) { + EvaluationContext.evaluate(ds.name(), dataset); + continue; + } + + StreamingQuery query = startQuery(dataset, checkpointBaseDir, leafIndex++, options); + boolean alreadyStopped; + synchronized (lock) { + queries.add(query); + alreadyStopped = stopped; + } + if (alreadyStopped) { + stopQuery(query); + } + } + + List<StreamingQuery> toAwait; + synchronized (lock) { + toAwait = new ArrayList<>(queries); + } + awaitTermination(toAwait); + } finally { + if (idleStopListener != null) { + getSparkSession().streams().removeListener(idleStopListener); + } + } + } + + /** + * Stops all queries started by {@link #evaluate()}. This method is idempotent and thread safe. + */ + @Override + public void stop() { + List<StreamingQuery> toStop; + synchronized (lock) { + if (stopped) { + return; + } + stopped = true; + toStop = new ArrayList<>(queries); + } + for (StreamingQuery query : toStop) { + stopQuery(query); + } + } + + private StreamingQuery startQuery( + Dataset<?> dataset, + String checkpointBaseDir, + int leafIndex, + SparkStructuredStreamingPipelineOptions options) { + try { + return dataset + .writeStream() + .format("noop") + .outputMode("append") + .option("checkpointLocation", 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<StreamingQuery> toAwait) { + List<StreamingQuery> active = new ArrayList<>(toAwait); + while (!active.isEmpty()) { + Iterator<StreamingQuery> iterator = active.iterator(); + while (iterator.hasNext()) { + StreamingQuery query = iterator.next(); + try { + if (query.awaitTermination(AWAIT_POLL_TIMEOUT_MILLIS)) { + iterator.remove(); + } + } catch (StreamingQueryException e) { + LOG.error("Streaming query {} terminated with an exception.", query.id(), e); + stop(); + throw new RuntimeException(e); + } + } + } + } + + /** Stops a single query if active. */ + private void stopQuery(StreamingQuery query) { + try { + if (query.isActive()) { + query.stop(); + } + } catch (TimeoutException | RuntimeException e) { + LOG.warn( + "Error while stopping streaming query {}: {}", + query.id(), + String.valueOf(e.getMessage())); Review Comment: Done, the exception is logged. ########## 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. + * + * <p>Leaf {@code i} checkpoints under {@code <checkpointDir>/i}, in pipeline graph order. A changed + * pipeline needs a new checkpoint directory, as with any Spark streaming query. + */ +@Internal +public class StreamingEvaluationContext extends EvaluationContext { + private static final Logger LOG = LoggerFactory.getLogger(StreamingEvaluationContext.class); + + private static final long AWAIT_POLL_TIMEOUT_MILLIS = 100; + + private final SparkStructuredStreamingPipelineOptions options; + + // Guards queries and stopped. + private final Object lock = new Object(); + private final List<StreamingQuery> queries = new ArrayList<>(); + private boolean stopped = false; + + StreamingEvaluationContext( + Collection<? extends NamedDataset<?>> leaves, + SparkSession session, + SparkCommonPipelineOptions options) { + super(leaves, session); + this.options = options.as(SparkStructuredStreamingPipelineOptions.class); + } + + /** Starts one streaming query per leaf dataset and blocks until all queries terminate. */ + @Override + public void evaluate() { + String checkpointBaseDir = options.getCheckpointDir(); + checkArgument( + checkpointBaseDir != null && !checkpointBaseDir.isEmpty(), + "checkpointDir must be set for a streaming pipeline"); + int idleStopThreshold = options.getStreamingStopAfterIdleBatches(); + + StreamingQueryListener idleStopListener = null; + if (idleStopThreshold >= 0) { + idleStopListener = new IdleStopListener(idleStopThreshold); + getSparkSession().streams().addListener(idleStopListener); + } + + try { + int leafIndex = 0; + for (NamedDataset<?> ds : leaves()) { + Dataset<?> dataset = ds.dataset(); + if (dataset == null) { + continue; + } + synchronized (lock) { + if (stopped) { + break; + } + } + if (!dataset.isStreaming()) { + EvaluationContext.evaluate(ds.name(), dataset); + continue; + } + + StreamingQuery query = startQuery(dataset, checkpointBaseDir, leafIndex++, options); + boolean alreadyStopped; + synchronized (lock) { + queries.add(query); + alreadyStopped = stopped; + } + if (alreadyStopped) { + stopQuery(query); + } + } + + List<StreamingQuery> toAwait; + synchronized (lock) { + toAwait = new ArrayList<>(queries); + } + awaitTermination(toAwait); + } finally { Review Comment: Done, a RuntimeException out of the try block stops every started query before it propagates. ########## 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. + * + * <p>Leaf {@code i} checkpoints under {@code <checkpointDir>/i}, in pipeline graph order. A changed + * pipeline needs a new checkpoint directory, as with any Spark streaming query. + */ +@Internal +public class StreamingEvaluationContext extends EvaluationContext { + private static final Logger LOG = LoggerFactory.getLogger(StreamingEvaluationContext.class); + + private static final long AWAIT_POLL_TIMEOUT_MILLIS = 100; + + private final SparkStructuredStreamingPipelineOptions options; + + // Guards queries and stopped. + private final Object lock = new Object(); + private final List<StreamingQuery> queries = new ArrayList<>(); + private boolean stopped = false; + + StreamingEvaluationContext( + Collection<? extends NamedDataset<?>> leaves, + SparkSession session, + SparkCommonPipelineOptions options) { + super(leaves, session); + this.options = options.as(SparkStructuredStreamingPipelineOptions.class); + } + + /** Starts one streaming query per leaf dataset and blocks until all queries terminate. */ + @Override + public void evaluate() { + String checkpointBaseDir = options.getCheckpointDir(); + checkArgument( + checkpointBaseDir != null && !checkpointBaseDir.isEmpty(), + "checkpointDir must be set for a streaming pipeline"); + int idleStopThreshold = options.getStreamingStopAfterIdleBatches(); + + StreamingQueryListener idleStopListener = null; + if (idleStopThreshold >= 0) { + idleStopListener = new IdleStopListener(idleStopThreshold); + getSparkSession().streams().addListener(idleStopListener); + } + + try { + int leafIndex = 0; + for (NamedDataset<?> ds : leaves()) { + Dataset<?> dataset = ds.dataset(); + if (dataset == null) { + continue; + } + synchronized (lock) { + if (stopped) { + break; + } + } + if (!dataset.isStreaming()) { + EvaluationContext.evaluate(ds.name(), dataset); + continue; + } + + StreamingQuery query = startQuery(dataset, checkpointBaseDir, leafIndex++, options); + boolean alreadyStopped; + synchronized (lock) { + queries.add(query); + alreadyStopped = stopped; + } + if (alreadyStopped) { + stopQuery(query); + } + } + + List<StreamingQuery> toAwait; + synchronized (lock) { + toAwait = new ArrayList<>(queries); + } + awaitTermination(toAwait); + } finally { + if (idleStopListener != null) { + getSparkSession().streams().removeListener(idleStopListener); + } + } + } + + /** + * Stops all queries started by {@link #evaluate()}. This method is idempotent and thread safe. + */ + @Override + public void stop() { + List<StreamingQuery> toStop; + synchronized (lock) { + if (stopped) { + return; + } + stopped = true; + toStop = new ArrayList<>(queries); + } + for (StreamingQuery query : toStop) { + stopQuery(query); + } + } + + private StreamingQuery startQuery( + Dataset<?> dataset, + String checkpointBaseDir, + int leafIndex, + SparkStructuredStreamingPipelineOptions options) { + try { + return dataset + .writeStream() + .format("noop") + .outputMode("append") + .option("checkpointLocation", checkpointBaseDir + "/" + leafIndex) Review Comment: Done, hadoop Path joins the checkpoint location. ########## 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. + * + * <p>Pipelines observe RUNNING, DONE once idle, CANCELLED on cancel, and FAILED on query failure. + */ +@RunWith(JUnit4.class) +@Category(StreamingTest.class) +public class StreamingPipelineLifecycleTest implements Serializable { + + /** Session shared across tests. */ + @ClassRule public static final SparkSessionRule SESSION = new SparkSessionRule(); + + @Rule public transient TemporaryFolder checkpointDir = new TemporaryFolder(); + + /** 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<String> collected = new HashSet<>(StreamingTestUtils.<String>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; Review Comment: Done, one awaitActiveQueries(IntPredicate, String) serves all three waits. ########## 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. + * + * <p>Pipelines observe RUNNING, DONE once idle, CANCELLED on cancel, and FAILED on query failure. + */ +@RunWith(JUnit4.class) +@Category(StreamingTest.class) +public class StreamingPipelineLifecycleTest implements Serializable { + + /** Session shared across tests. */ + @ClassRule public static final SparkSessionRule SESSION = new SparkSessionRule(); + + @Rule public transient TemporaryFolder checkpointDir = new TemporaryFolder(); + + /** 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<String> collected = new HashSet<>(StreamingTestUtils.<String>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) { Review Comment: Same, uses awaitActiveQueries. ########## 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. + * + * <p>Pipelines observe RUNNING, DONE once idle, CANCELLED on cancel, and FAILED on query failure. + */ +@RunWith(JUnit4.class) +@Category(StreamingTest.class) +public class StreamingPipelineLifecycleTest implements Serializable { + + /** Session shared across tests. */ + @ClassRule public static final SparkSessionRule SESSION = new SparkSessionRule(); + + @Rule public transient TemporaryFolder checkpointDir = new TemporaryFolder(); + + /** 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"); Review Comment: Removed here and in StatelessParDoStreamingTest, neither asserts on the source counters. The restart test keeps its @After because it does. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
