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..3529daaff313 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java @@ -0,0 +1,113 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.spark.structuredstreaming.translation; + +import java.util.Collection; +import org.apache.beam.runners.spark.SparkCommonPipelineOptions; +import org.apache.beam.runners.spark.structuredstreaming.translation.batch.PipelineTranslatorCommon; +import org.apache.beam.runners.spark.structuredstreaming.translation.streaming.ReadUnboundedTranslator; +import org.apache.beam.sdk.annotations.Internal; +import org.apache.beam.sdk.transforms.Combine; +import org.apache.beam.sdk.transforms.GroupByKey; +import org.apache.beam.sdk.transforms.Impulse; +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.reflect.DoFnSignature; +import org.apache.beam.sdk.transforms.reflect.DoFnSignatures; +import org.apache.beam.sdk.util.construction.SplittableParDo; +import org.apache.beam.sdk.values.PInput; +import org.apache.beam.sdk.values.POutput; +import org.apache.spark.sql.SparkSession; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * Pipeline translator for streaming pipelines on Spark 4. It extends the common registry to reuse + * the stateless single output ParDo, Window.Assign, Flatten and Reshuffle translators, which are + * safe on a streaming Dataset. Every other primitive fails at translation, the batch translators + * for them persist or collect the Dataset, which Spark rejects for streaming plans. + */ +@Internal +public class PipelineTranslatorStreaming extends PipelineTranslatorCommon { + + @SuppressWarnings("rawtypes") + private static final TransformTranslator READ_UNBOUNDED = new ReadUnboundedTranslator<>(); + + private static final String NOT_SUPPORTED = + " is not supported by the Spark 4 streaming runner yet, see" + + " https://github.com/apache/beam/issues/36841"; + + /** Returns a {@link TransformTranslator} for the given {@link PTransform} if known. */ + @Override + @Nullable + protected > + TransformTranslator getTransformTranslator(TransformT transform) { + + if (transform instanceof SplittableParDo.PrimitiveUnboundedRead) { + return READ_UNBOUNDED; + } + + if (transform instanceof SplittableParDo.PrimitiveBoundedRead) { + throw new UnsupportedOperationException( + "Bounded Read (Read.from(BoundedSource), Create with two or more elements)" + + NOT_SUPPORTED); + } + + if (transform instanceof Impulse) { + throw new UnsupportedOperationException( + "Impulse (Create with fewer than two elements, PAssert)" + NOT_SUPPORTED); + } + + if (transform instanceof GroupByKey) { + throw new UnsupportedOperationException("GroupByKey" + NOT_SUPPORTED); + } + + if (transform instanceof Combine.PerKey) { + throw new UnsupportedOperationException("Combine.perKey" + NOT_SUPPORTED); + } + + if (transform instanceof ParDo.MultiOutput) { + ParDo.MultiOutput parDo = (ParDo.MultiOutput) transform; + DoFnSignature signature = DoFnSignatures.signatureForDoFn(parDo.getFn()); + if (signature.usesState() || signature.usesTimers()) { + throw new UnsupportedOperationException( + "Stateful ParDo (" + signature.fnClass().getName() + ")" + NOT_SUPPORTED); + } + if (!parDo.getSideInputs().isEmpty()) { + throw new UnsupportedOperationException( + "ParDo with side inputs (" + signature.fnClass().getName() + ")" + NOT_SUPPORTED); + } + if (!parDo.getAdditionalOutputTags().getAll().isEmpty()) { + throw new UnsupportedOperationException( + "ParDo with additional outputs (" + + signature.fnClass().getName() + + ")" + + NOT_SUPPORTED); + } + } + + return super.getTransformTranslator(transform); + } + + @Override + protected EvaluationContext createEvaluationContext( + Collection> 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..ff7839bddb20 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/StreamingEvaluationContext.java @@ -0,0 +1,261 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.spark.structuredstreaming.translation; + +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicInteger; +import org.apache.beam.runners.spark.SparkCommonPipelineOptions; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +import org.apache.beam.sdk.annotations.Internal; +import org.apache.hadoop.fs.Path; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.streaming.StreamingQuery; +import org.apache.spark.sql.streaming.StreamingQueryException; +import org.apache.spark.sql.streaming.StreamingQueryListener; +import org.apache.spark.sql.streaming.Trigger; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Starts one Spark Structured Streaming query per leaf dataset and blocks until all of them reach a + * terminal state. Queries end through {@link #stop()} or the idle stop listener. + * + *

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); + } catch (RuntimeException e) { + stop(); + throw e; + } finally { + if (idleStopListener != null) { + getSparkSession().streams().removeListener(idleStopListener); + } + } + } + + /** + * Stops all queries started by {@link #evaluate()}. This method is idempotent and thread safe. + */ + @Override + public void stop() { + List toStop; + synchronized (lock) { + if (stopped) { + return; + } + stopped = true; + toStop = new ArrayList<>(queries); + } + for (StreamingQuery query : toStop) { + stopQuery(query); + } + } + + private StreamingQuery startQuery( + Dataset dataset, + String checkpointBaseDir, + int leafIndex, + SparkStructuredStreamingPipelineOptions options) { + try { + return dataset + .writeStream() + .format("noop") + .outputMode("append") + .option( + "checkpointLocation", + new Path(checkpointBaseDir, Integer.toString(leafIndex)).toString()) + .trigger(Trigger.ProcessingTime(options.getMaxBatchDurationMillis())) + .start(); + } catch (TimeoutException e) { + throw new RuntimeException( + "Failed to start streaming query for leaf dataset index " + leafIndex, e); + } + } + + /** Blocks until every query in toAwait has terminated. Sibling queries stop on failure. */ + private void awaitTermination(List 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("Failed to stop streaming query {}.", query.id(), e); + } + } + + private void stopQueryById(UUID id) { + StreamingQuery match = null; + synchronized (lock) { + for (StreamingQuery query : queries) { + if (query.id().equals(id)) { + match = query; + break; + } + } + } + if (match != null) { + stopQuery(match); + } + } + + /** + * Stops a query after {@code threshold} consecutive triggers without input rows. A trigger + * without data is reported as a progress event with zero rows when the source offset moved and as + * an idle event otherwise, both count. + */ + private final class IdleStopListener extends StreamingQueryListener { + private final int threshold; + private final Map 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..d41f5e774227 --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatelessParDoStreamingTest.java @@ -0,0 +1,110 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.spark.structuredstreaming.translation.streaming; + +import static org.junit.Assert.assertEquals; + +import java.io.Serializable; +import java.util.HashSet; +import java.util.Set; +import org.apache.beam.runners.spark.StreamingTest; +import org.apache.beam.runners.spark.structuredstreaming.SparkSessionRule; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +import org.apache.beam.runners.spark.structuredstreaming.io.streaming.TestUnboundedSource; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.io.Read; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.Flatten; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionList; +import org.junit.ClassRule; +import org.junit.Rule; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Baseline streaming pipeline tests for stateless ParDo and Flatten. */ +@RunWith(JUnit4.class) +@Category(StreamingTest.class) +public class StatelessParDoStreamingTest implements Serializable { + + @ClassRule public static final SparkSessionRule SESSION = new SparkSessionRule(); + + @Rule public transient TemporaryFolder checkpointDir = new TemporaryFolder(); + + private static class PassThroughFn extends DoFn { + @ProcessElement + public void process(@Element String element, OutputReceiver out) { + out.output(element); + } + } + + @Test + public void everyElementPassesThrough() throws Exception { + String tag = "stateless-pardo"; + String collectorId = StreamingTestUtils.newCollectorId(tag); + StreamingTestUtils.clear(collectorId); + + SparkStructuredStreamingPipelineOptions options = + StreamingTestUtils.streamingOptions(checkpointDir); + Pipeline pipeline = Pipeline.create(options); + + pipeline + .apply("ReadUnbounded", Read.from(new TestUnboundedSource(tag, 1, 10))) + .apply("PassThrough", ParDo.of(new PassThroughFn())) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + + PipelineResult result = StreamingTestUtils.run(pipeline); + + Set collected = StreamingTestUtils.collected(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 = StreamingTestUtils.collected(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..c7ad2fa62956 --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingCheckpointRestartTest.java @@ -0,0 +1,118 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.spark.structuredstreaming.translation.streaming; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +import java.io.File; +import java.io.Serializable; +import java.util.Collections; +import java.util.Set; +import org.apache.beam.runners.spark.StreamingTest; +import org.apache.beam.runners.spark.structuredstreaming.SparkSessionRule; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +import org.apache.beam.runners.spark.structuredstreaming.io.streaming.BeamReaderCache; +import org.apache.beam.runners.spark.structuredstreaming.io.streaming.TestUnboundedSource; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.io.Read; +import org.apache.beam.sdk.transforms.ParDo; +import org.junit.After; +import org.junit.ClassRule; +import org.junit.Rule; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * Verifies restarted pipelines resume from durable checkpoint marks. + * + *

Two pipelines run sequentially against the same checkpoint directory. Wiping the reader cache + * in between forces the second run to restore from the durable marks. + */ +@RunWith(JUnit4.class) +@Category(StreamingTest.class) +public class StreamingCheckpointRestartTest implements Serializable { + + @ClassRule public static final SparkSessionRule SESSION = new SparkSessionRule(); + + @Rule public transient TemporaryFolder checkpointDir = new TemporaryFolder(); + + private static final int ELEMENT_COUNT = 10; + private static final String TAG = "checkpoint-restart"; + + @After + public void tearDown() { + TestUnboundedSource.forget(TAG); + } + + @Test + public void restartedPipelineResumesFromDurableCheckpointMarks() throws Exception { + String checkpointPath = checkpointDir.newFolder("checkpoint").getAbsolutePath(); + + String collectorA = StreamingTestUtils.newCollectorId("checkpoint-restart-a"); + String collectorB = StreamingTestUtils.newCollectorId("checkpoint-restart-b"); + StreamingTestUtils.clear(collectorA); + StreamingTestUtils.clear(collectorB); + + // First run reads all elements from a fresh checkpoint directory. + runPipeline(checkpointPath, collectorA, TAG); + + Set collectedA = StreamingTestUtils.collected(collectorA); + assertEquals( + "first run must read every element", + TestUnboundedSource.elements(TAG, 1, ELEMENT_COUNT), + collectedA); + + // The source checkpoint lives under the location the translator handed to Spark. + File sourceRoot = new File(new File(checkpointPath, "0"), "sources/0"); + assertTrue("expected source checkpoint directory " + sourceRoot, sourceRoot.isDirectory()); + + int createdBeforeSecondRun = TestUnboundedSource.created(TAG); + + // Invalidate cached readers and marks to force restore from durable files. + BeamReaderCache.invalidateAll(); + + // Second run resumes against the same checkpoint directory. + runPipeline(checkpointPath, collectorB, TAG); + + Set collectedB = StreamingTestUtils.collected(collectorB); + assertTrue( + "second run must not re-emit elements the first run committed", + Collections.disjoint(collectedA, collectedB)); + assertTrue( + "readers must be recreated during the second run", + TestUnboundedSource.created(TAG) > createdBeforeSecondRun); + } + + private PipelineResult runPipeline(String checkpointPath, String collectorId, String tag) { + SparkStructuredStreamingPipelineOptions options = + StreamingTestUtils.streamingOptions(checkpointPath); + // One element per split per micro batch to spread elements across batches. + options.setMaxRecordsPerBatch(1L); + + Pipeline pipeline = Pipeline.create(options); + pipeline + .apply("ReadUnbounded", Read.from(new TestUnboundedSource(tag, 1, ELEMENT_COUNT))) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + return StreamingTestUtils.run(pipeline); + } +} diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java new file mode 100644 index 000000000000..76462c38fe10 --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java @@ -0,0 +1,171 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.spark.structuredstreaming.translation.streaming; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +import java.io.Serializable; +import java.util.Set; +import java.util.function.IntPredicate; +import org.apache.beam.runners.spark.StreamingTest; +import org.apache.beam.runners.spark.structuredstreaming.SparkSessionRule; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +import org.apache.beam.runners.spark.structuredstreaming.io.streaming.TestUnboundedSource; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.io.Read; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.junit.ClassRule; +import org.junit.Rule; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * State transitions for a streaming pipeline. + * + *

Pipelines observe RUNNING, DONE once idle, CANCELLED on cancel, and FAILED on query failure. + */ +@RunWith(JUnit4.class) +@Category(StreamingTest.class) +public class StreamingPipelineLifecycleTest implements Serializable { + + /** Session shared across tests. */ + @ClassRule public static final SparkSessionRule SESSION = new SparkSessionRule(); + + @Rule public transient TemporaryFolder checkpointDir = new TemporaryFolder(); + + /** Upper bound for a query to start or stop. */ + private static final long QUERY_START_TIMEOUT_MILLIS = 60_000L; + + /** Polls the active query count of the shared session until {@code condition} holds. */ + private static void awaitActiveQueries(IntPredicate condition, String failure) + throws InterruptedException { + long deadline = System.currentTimeMillis() + QUERY_START_TIMEOUT_MILLIS; + while (!condition.test(SESSION.getSession().streams().active().length)) { + assertTrue( + failure + " within " + QUERY_START_TIMEOUT_MILLIS + " ms", + System.currentTimeMillis() < deadline); + Thread.sleep(50L); + } + } + + @Test + public void idlePipelineGoesFromRunningToDoneOnceIdle() throws Exception { + String tag = "lifecycle-done"; + String collectorId = StreamingTestUtils.newCollectorId(tag); + StreamingTestUtils.clear(collectorId); + + SparkStructuredStreamingPipelineOptions options = + StreamingTestUtils.streamingOptions(checkpointDir); + Pipeline pipeline = Pipeline.create(options); + + pipeline + .apply("ReadUnbounded", Read.from(new TestUnboundedSource(tag, 1, 10))) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + + PipelineResult result = pipeline.run(); + assertEquals(PipelineResult.State.RUNNING, result.getState()); + + PipelineResult.State finalState = StreamingTestUtils.waitUntilFinish(result); + assertEquals(PipelineResult.State.DONE, finalState); + assertEquals(PipelineResult.State.DONE, result.getState()); + + Set collected = StreamingTestUtils.collected(collectorId); + assertEquals(TestUnboundedSource.elements(tag, 1, 10), collected); + } + + @Test + public void cancelStopsTheQueryAndReportsCancelled() throws Exception { + String tag = "lifecycle-cancel"; + String collectorId = StreamingTestUtils.newCollectorId(tag); + StreamingTestUtils.clear(collectorId); + + SparkStructuredStreamingPipelineOptions options = + StreamingTestUtils.streamingOptions(checkpointDir); + // Idle stop disabled so the query stops only from explicit cancel. + options.setStreamingStopAfterIdleBatches(-1); + Pipeline pipeline = Pipeline.create(options); + + pipeline + .apply("ReadUnbounded", Read.from(new TestUnboundedSource(tag, 1, 10))) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + + PipelineResult result = pipeline.run(); + assertEquals(PipelineResult.State.RUNNING, result.getState()); + + awaitActiveQueries(count -> count > 0, "no streaming query started"); + + PipelineResult.State cancelledState = result.cancel(); + assertEquals(PipelineResult.State.CANCELLED, cancelledState); + assertEquals(PipelineResult.State.CANCELLED, result.getState()); + + awaitActiveQueries(count -> count == 0, "the streaming query did not stop after cancel"); + } + + /** A failure in any leaf query surfaces through waitUntilFinish. */ + @Test + public void failingLeafQueryFailsThePipelineAndStopsHealthySibling() throws Exception { + String healthyTag = "lifecycle-healthy"; + String poisonTag = "lifecycle-poison"; + String collectorId = StreamingTestUtils.newCollectorId(healthyTag); + StreamingTestUtils.clear(collectorId); + + SparkStructuredStreamingPipelineOptions options = + StreamingTestUtils.streamingOptions(checkpointDir); + // Idle stop disabled so the healthy query stops only when the sibling failure stops it. + options.setStreamingStopAfterIdleBatches(-1); + Pipeline pipeline = Pipeline.create(options); + + pipeline + .apply("ReadHealthy", Read.from(new TestUnboundedSource(healthyTag, 1, 10))) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + pipeline + .apply("ReadPoisoned", Read.from(new TestUnboundedSource(poisonTag, 1, 10))) + .apply("Throw", ParDo.of(new ThrowOnElementDoFn(5))); + + PipelineResult result = pipeline.run(); + + assertThrows(RuntimeException.class, () -> StreamingTestUtils.waitUntilFinish(result)); + assertEquals(PipelineResult.State.FAILED, result.getState()); + + awaitActiveQueries(count -> count == 0, "a sibling query did not stop after the failure"); + } + + /** Throws on one specific element index, passes every other element through. */ + private static final class ThrowOnElementDoFn extends DoFn { + 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..3d7841154150 --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTestUtils.java @@ -0,0 +1,156 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.spark.structuredstreaming.translation.streaming; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingRunner; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; +import org.joda.time.Duration; +import org.junit.rules.TemporaryFolder; + +/** + * Shared test utilities for the Spark 4 streaming translators. + * + *

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 Set collected(String collectorId) { + List values = COLLECTORS.get(collectorId); + if (values == null) { + return Collections.emptySet(); + } + synchronized (values) { + return (Set) new HashSet<>(values); + } + } + + /** Discards everything collected so far under {@code collectorId}. */ + public static void clear(String collectorId) { + COLLECTORS.remove(collectorId); + } + + /** Collector id that will not collide with other tests or runs. */ + public static String newCollectorId(String prefix) { + return prefix + "-" + UUID.randomUUID(); + } + + /** Upper bound on the wall clock time one streaming pipeline may take. */ + public static final Duration FINISH_TIMEOUT = Duration.standardMinutes(5); + + /** + * Runs {@code pipeline} and returns on a terminal state. Cancels the pipeline and fails after + * {@link #FINISH_TIMEOUT}. + */ + public static PipelineResult run(Pipeline pipeline) { + PipelineResult result = pipeline.run(); + waitUntilFinish(result); + return result; + } + + /** + * Waits at most {@link #FINISH_TIMEOUT} for {@code result} to reach a terminal state. Fails and + * cancels if still running at the deadline. + */ + public static PipelineResult.State waitUntilFinish(PipelineResult result) { + PipelineResult.State state = result.waitUntilFinish(FINISH_TIMEOUT); + if (state == null || !state.isTerminal()) { + try { + result.cancel(); + } catch (IOException | RuntimeException e) { + // Best effort cancellation. + } + throw new AssertionError( + "pipeline did not finish within " + FINISH_TIMEOUT + ", last state " + state); + } + return state; + } + + /** + * Streaming options for tests: runner on the active session, streaming mode, stop after 3 idle + * batches, 200 ms trigger, the given checkpoint directory. Test mode is off, {@code run()} + * returns at once and tests wait through {@link #run} or {@link #waitUntilFinish}. + */ + public static SparkStructuredStreamingPipelineOptions streamingOptions( + TemporaryFolder checkpointDir) throws IOException { + return streamingOptions(checkpointDir.newFolder("checkpoint").getAbsolutePath()); + } + + /** Same as {@link #streamingOptions(TemporaryFolder)} with an explicit checkpoint path. */ + public static SparkStructuredStreamingPipelineOptions streamingOptions(String checkpointPath) { + SparkStructuredStreamingPipelineOptions options = + PipelineOptionsFactory.as(SparkStructuredStreamingPipelineOptions.class); + options.setRunner(SparkStructuredStreamingRunner.class); + options.setUseActiveSparkSession(true); + options.setTestMode(false); + options.setStreaming(true); + options.setStreamingStopAfterIdleBatches(3); + options.setMaxBatchDurationMillis(200); + options.setCheckpointDir(checkpointPath); + return options; + } +} diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/ParDoTranslatorBatch.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/ParDoTranslatorBatch.java index 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..899e0443a94f --- /dev/null +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/PipelineTranslatorCommon.java @@ -0,0 +1,109 @@ +/* + * 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.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}. *