diff --git a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark4.json b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark4.json index e0266d62f2e0..f1ba03a243ee 100644 --- a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark4.json +++ b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark4.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run", - "modification": 4 + "modification": 5 } diff --git a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json index 1b832d7f19cb..65728e1581dc 100644 --- a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json +++ b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json @@ -12,5 +12,5 @@ "https://github.com/apache/beam/pull/40103": "noting that PR #40103 should run this test", "https://github.com/apache/beam/issues/40427": "Spark ValidatesRunner tests run serially", "https://github.com/apache/beam/pull/40281": "noting that PR #40281 should run this test", - "https://github.com/apache/beam/issues/36841": "Spark 4 streaming GroupByKey" + "https://github.com/apache/beam/issues/36841": "Spark 4 streaming Impulse and PAssert" } diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java index e961ff2f3807..8348d04d289c 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java @@ -32,6 +32,7 @@ import org.apache.beam.sdk.io.UnboundedSource.CheckpointMark; import org.apache.beam.sdk.io.UnboundedSource.UnboundedReader; import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; import org.checkerframework.checker.nullness.qual.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -231,6 +232,11 @@ public static final class CachedReader implements Closeable { private final long idleTimeoutMillis; private final LongSupplier committedEpoch; private boolean started; + private boolean exhausted; + private boolean positionExhausted; + private boolean doneEligible; + private boolean doneWritten; + private boolean sentinelEmitted; private boolean inBatch; private boolean moved; private long positionEpoch; @@ -258,11 +264,31 @@ public UnboundedReader reader() { public synchronized boolean startOrAdvance() throws IOException { moved = true; - if (!started) { - started = true; - return reader.start(); + boolean hasElement = started ? reader.advance() : reader.start(); + started = true; + if (!hasElement && !reader.getWatermark().isBefore(GlobalWindow.INSTANCE.maxTimestamp())) { + exhausted = true; } - return reader.advance(); + return hasElement; + } + + /** + * Claims the one time end of stream sentinel for this reader once the driver has seen every + * split done. + */ + public synchronized boolean claimSentinel(boolean emitSentinel) { + if (!emitSentinel || sentinelEmitted) { + return false; + } + sentinelEmitted = true; + return true; + } + + /** + * Whether a committed batch exhausted this reader and its done marker has not been written yet. + */ + synchronized boolean shouldWriteDone() { + return doneEligible && !doneWritten; } /** @@ -288,6 +314,7 @@ synchronized boolean beginBatch(long epoch) { if (positionEpoch != epoch || moved) { return false; } + doneEligible |= positionExhausted; inBatch = true; lastUsedMillis = System.currentTimeMillis(); return true; @@ -298,6 +325,8 @@ synchronized void endBatch(long endEpoch, @Nullable CheckpointMark mark, byte[] positionEpoch = endEpoch; positionMark = codedMark; pendingMark = mark; + positionExhausted = exhausted; + doneWritten |= doneEligible; moved = false; inBatch = false; lastUsedMillis = System.currentTimeMillis(); diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamSourceCheckpoint.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamSourceCheckpoint.java index d667976edcec..8b58a0e51070 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamSourceCheckpoint.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamSourceCheckpoint.java @@ -42,7 +42,8 @@ * *

{@code /splits} pins the split list, written once by the driver. {@code * /marks//} holds the coded checkpoint mark of a split at the end of the - * batch ending at that epoch. The epoch Spark last committed is read from Spark's own {@code + * batch ending at that epoch. {@code /done/} marks a split whose reader reached + * the end of the global window. The epoch Spark last committed is read from Spark's own {@code * commits} and {@code offsets} logs two levels up. All IO goes through Spark's {@link * CheckpointFileManager}, writes are atomic renames. */ @@ -52,6 +53,7 @@ public final class BeamSourceCheckpoint { private static final String SPLITS_FILE = "splits"; private static final String MARKS_DIR = "marks"; + private static final String DONE_DIR = "done"; private static final String SPARK_COMMITS_DIR = "commits"; private static final String SPARK_OFFSETS_DIR = "offsets"; private static final String SERIALIZED_VOID_OFFSET = "-"; @@ -64,6 +66,7 @@ public final class BeamSourceCheckpoint { private final Path root; private final Path splitsPath; private final Path marksRoot; + private final Path doneRoot; /** Every mark epoch strictly below the value is known to be deleted, -1 for unknown. */ private volatile long purgeFloor = -1L; @@ -74,6 +77,7 @@ public BeamSourceCheckpoint(String checkpointLocation, Configuration hadoopConf) this.fm = CheckpointFileManager.create(root, hadoopConf); this.splitsPath = new Path(root, SPLITS_FILE); this.marksRoot = new Path(root, MARKS_DIR); + this.doneRoot = new Path(root, DONE_DIR); } public String location() { @@ -129,6 +133,32 @@ public void writeMark(int splitId, long epoch, byte[] codedMark) throws IOExcept return read(path); } + /** Writes the one time done marker of {@code splitId}, idempotent. */ + public void writeDone(int splitId) throws IOException { + Path path = new Path(doneRoot, Integer.toString(splitId)); + try { + write(path, new byte[0], true); + } catch (FileNotFoundException e) { + fm.mkdirs(doneRoot); + write(path, new byte[0], true); + } + } + + /** Whether every split in {@code [0, numSplits)} has a marker under {@code done/}. */ + public boolean allSplitsDone(int numSplits) throws IOException { + if (!fm.exists(doneRoot)) { + return false; + } + int count = 0; + for (FileStatus status : fm.list(doneRoot)) { + long splitId = parseEpoch(status.getPath().getName()); + if (splitId >= 0 && splitId < numSplits && ++count == numSplits) { + return true; + } + } + return false; + } + /** * The end epoch of this source in the last batch Spark committed, or -1 if there is none or the * logs cannot be read. The location is {@code /sources/}, the batch id is the diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/UnboundedSourceDataset.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/UnboundedSourceDataset.java index 33106824c107..3ef35762748a 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/UnboundedSourceDataset.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/UnboundedSourceDataset.java @@ -36,6 +36,7 @@ import org.apache.beam.sdk.io.UnboundedSource; import org.apache.beam.sdk.io.UnboundedSource.CheckpointMark; import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.util.BackOff; import org.apache.beam.sdk.util.BackOffUtils; import org.apache.beam.sdk.util.FluentBackoff; @@ -86,7 +87,9 @@ * TIMESTAMP} holding the event timestamp of that element. * *

The event time watermark is declared here and only here. Spark 4 rejects a second {@code - * withWatermark} further down the plan, so downstream translators must never call it again. + * withWatermark} further down the plan, so downstream translators must never call it again. Once + * every split of the source is done, each reader emits one empty payload row at {@link + * BoundedWindow#TIMESTAMP_MAX_VALUE} to advance the watermark, which is dropped at decode. */ public final class UnboundedSourceDataset { @@ -305,6 +308,7 @@ static class BeamMicroBatchStream implements MicroBatchStream { private final AtomicLong purgeRequested = new AtomicLong(); private long epoch; + private boolean emitSentinel; private @Nullable List> splits; BeamMicroBatchStream(BeamTable table, String checkpointLocation) { @@ -343,6 +347,13 @@ public InputPartition[] planInputPartitions(Offset start, Offset end) { LOG.warn( "Failed to prepare mark directory of epoch {} at {}.", endEpoch, checkpointLocation, e); } + if (!emitSentinel) { + try { + emitSentinel = checkpoint.allSplitsDone(pinned.size()); + } catch (IOException e) { + LOG.warn("Failed to list done splits at {}.", checkpointLocation, e); + } + } long[] quotas = splitQuotas(table.maxRecordsPerBatch, pinned.size(), endEpoch); InputPartition[] partitions = new InputPartition[pinned.size()]; for (int i = 0; i < pinned.size(); i++) { @@ -358,7 +369,8 @@ public InputPartition[] planInputPartitions(Offset start, Offset end) { endEpoch, quotas[i], table.maxBatchDurationMillis, - table.readerIdleTimeoutMillis); + table.readerIdleTimeoutMillis, + emitSentinel); } return partitions; } @@ -495,6 +507,7 @@ static final class BeamInputPartition implements InputPartition { final long maxBatchDurationMillis; final long readerIdleTimeoutMillis; + final boolean emitSentinel; BeamInputPartition( UnboundedSource split, @@ -507,7 +520,8 @@ static final class BeamInputPartition implements InputPartition { long endEpoch, long maxRecords, long maxBatchDurationMillis, - long readerIdleTimeoutMillis) { + long readerIdleTimeoutMillis, + boolean emitSentinel) { this.split = split; this.coder = coder; this.options = options; @@ -519,6 +533,7 @@ static final class BeamInputPartition implements InputPartition { this.maxRecords = maxRecords; this.maxBatchDurationMillis = maxBatchDurationMillis; this.readerIdleTimeoutMillis = readerIdleTimeoutMillis; + this.emitSentinel = emitSentinel; } @Override @@ -573,6 +588,7 @@ static final class BeamPartitionReader implements PartitionReader implements PartitionReader getTransformTranslator(TransformT tra 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) { return GROUP_BY_KEY; } @@ -175,7 +164,8 @@ private static void checkNoProcessingTimeTimers(DoFn doFn, DoFnSignature s private static void checkSupportedGroupByKey(PCollection input) { WindowingStrategy strategy = input.getWindowingStrategy(); checkNotMergingWindows(strategy, "GroupByKey"); - if (!(strategy.getTrigger() instanceof DefaultTrigger)) { + if (!(strategy.getTrigger() instanceof DefaultTrigger + || strategy.getTrigger() instanceof NeverTrigger)) { throw new UnsupportedOperationException( "Custom triggers (" + strategy.getTrigger() + ")" + NOT_SUPPORTED); } diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/StreamingEvaluationContext.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/StreamingEvaluationContext.java index 2ca54a1e9481..b5057cd260ac 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/StreamingEvaluationContext.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/StreamingEvaluationContext.java @@ -30,6 +30,7 @@ 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.Pipeline; import org.apache.beam.sdk.annotations.Internal; import org.apache.hadoop.fs.Path; import org.apache.spark.sql.Dataset; @@ -179,7 +180,7 @@ private void awaitTermination(List toAwait) { } catch (StreamingQueryException e) { LOG.error("Streaming query {} terminated with an exception.", query.id(), e); stop(); - throw new RuntimeException(e); + throw new Pipeline.PipelineExecutionException(e.getCause() != null ? e.getCause() : e); } } } 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 index eb1421b45540..2281bbfdec31 100644 --- 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 @@ -19,6 +19,8 @@ import static org.apache.spark.sql.functions.col; +import java.util.Collections; +import java.util.Iterator; import org.apache.beam.runners.spark.coders.CoderHelpers; import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; import org.apache.beam.runners.spark.structuredstreaming.io.streaming.UnboundedSourceDataset; @@ -31,7 +33,7 @@ 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.api.java.function.FlatMapFunction; import org.apache.spark.sql.Dataset; import org.apache.spark.sql.Encoder; import org.apache.spark.sql.Encoders; @@ -76,13 +78,16 @@ protected void translate(SplittableParDo.PrimitiveUnboundedRead transform, Co Dataset> dataset = rows.select(col(UnboundedSourceDataset.COL_PAYLOAD)) .as(Encoders.BINARY()) - .map(new DecodePayload<>(payloadCoder), encoder); + .flatMap(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> { + /** + * Decodes the binary payload column back into a Beam {@code WindowedValue}, dropping the empty + * end of stream sentinel payload. + */ + private static final class DecodePayload implements FlatMapFunction> { private final Coder> coder; DecodePayload(Coder> coder) { @@ -90,8 +95,11 @@ private static final class DecodePayload implements MapFunction call(byte[] payload) { - return CoderHelpers.fromByteArray(payload, coder); + public Iterator> call(byte[] payload) { + if (payload.length == 0) { + return Collections.emptyIterator(); + } + return Collections.singleton(CoderHelpers.fromByteArray(payload, coder)).iterator(); } } } 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 559deb808036..07c823068c7d 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 @@ -71,7 +71,6 @@ import org.apache.spark.sql.Encoders; import org.apache.spark.sql.Row; import org.apache.spark.sql.SparkSession; -import org.apache.spark.sql.catalyst.InternalRow; import org.apache.spark.sql.catalyst.plans.logical.EventTimeWatermark; import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan; import org.apache.spark.sql.connector.read.InputPartition; @@ -547,7 +546,13 @@ private static StreamingQuery start(Dataset dataset, String queryName, File (batch, batchId) -> { List> target = BATCHES.get(queryName); if (target != null) { - target.add(batch.collectAsList()); + List rows = new ArrayList<>(); + for (Row row : batch.collectAsList()) { + if (((byte[]) row.getAs(COL_PAYLOAD)).length > 0) { + rows.add(row); + } + } + target.add(rows); } }) .queryName(queryName) @@ -766,7 +771,8 @@ private BeamInputPartition partition(String location, long start, long e end, 3L, 30_000L, - 600_000L); + 600_000L, + false); } private static List readBatch(BeamInputPartition partition) throws IOException { @@ -778,8 +784,10 @@ private static List readBatch(BeamInputPartition partition) thro private static void drainInto(BeamPartitionReader reader, List values) throws IOException { while (reader.next()) { - InternalRow row = reader.get(); - values.add(decode(row.getBinary(0))); + byte[] payload = reader.get().getBinary(0); + if (payload.length > 0) { + values.add(decode(payload)); + } } reader.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 index 98f62501f378..8bc257d18729 100644 --- 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 @@ -41,8 +41,9 @@ /** * 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. + * formatted as tag, shard, index. Watermark advances to TIMESTAMP_MAX_VALUE after exhaustion, while + * an {@link #extend} below the declared count holds the watermark. Marks record finalized positions + * under tag and shard. */ public final class TestUnboundedSource extends UnboundedSource { private static final long serialVersionUID = 1L; @@ -62,20 +63,27 @@ public final class TestUnboundedSource extends UnboundedSource elements(String tag, int shards, int count) { @@ -88,6 +96,16 @@ public static Set elements(String tag, int shards, int count) { return elements; } + public static Set elements(String tag, int[] shardCounts) { + Set elements = new HashSet<>(); + for (int shard = 0; shard < shardCounts.length; shard++) { + for (int index = 0; index < shardCounts[shard]; index++) { + elements.add(element(tag, shard, index)); + } + } + return elements; + } + public static String element(String tag, int shard, int index) { return tag + "-" + shard + "-" + index; } @@ -133,7 +151,8 @@ public List split(int desiredNumSplits, PipelineOptions opt } List splits = new ArrayList<>(); for (int i = 0; i < shards; i++) { - splits.add(new TestUnboundedSource(tag, i, shards, perShard)); + int count = shardCounts != null ? shardCounts[i] : perShard; + splits.add(new TestUnboundedSource(tag, i, shards, count, null)); } return splits; } @@ -200,6 +219,7 @@ private static final class Reader extends UnboundedReader { private final TestUnboundedSource source; private int next; private int current = -1; + private boolean exhausted; Reader(TestUnboundedSource source, int next) { this.source = source; @@ -222,6 +242,7 @@ public boolean advance() { current = next++; return true; } + exhausted = next >= source.perShard; return false; } @@ -238,11 +259,15 @@ public Instant getCurrentTimestamp() throws NoSuchElementException { if (current < 0) { throw new NoSuchElementException(); } - return new Instant(BASE_MILLIS + (source.shard * limit() + current) * INTERVAL_MILLIS); + return new Instant( + BASE_MILLIS + (source.shard * source.perShard + current) * INTERVAL_MILLIS); } @Override public Instant getWatermark() { + if (exhausted) { + return BoundedWindow.TIMESTAMP_MAX_VALUE; + } return current < 0 ? BoundedWindow.TIMESTAMP_MIN_VALUE : getCurrentTimestamp(); } diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreamingTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreamingTest.java index a5afa7c1a4a6..d18f069cfbec 100644 --- a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreamingTest.java +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreamingTest.java @@ -34,7 +34,6 @@ import org.apache.beam.sdk.state.TimerSpec; import org.apache.beam.sdk.state.TimerSpecs; import org.apache.beam.sdk.transforms.Count; -import org.apache.beam.sdk.transforms.Create; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.GroupByKey; import org.apache.beam.sdk.transforms.ParDo; @@ -87,22 +86,6 @@ private static void assertUnsupported(Pipeline pipeline, String expected) { throw new AssertionError("missing " + expected); } - @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 rejectsUnsupportedStatefulFeatures() throws Exception { Object[][] cases = { diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/PAssertStreamingTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/PAssertStreamingTest.java new file mode 100644 index 000000000000..79a836f1c02e --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/PAssertStreamingTest.java @@ -0,0 +1,149 @@ +/* + * 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.assertThrows; +import static org.junit.Assert.assertTrue; + +import java.io.IOException; +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.sdk.Pipeline; +import org.apache.beam.sdk.io.Read; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.windowing.FixedWindows; +import org.apache.beam.sdk.transforms.windowing.Window; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Throwables; +import org.joda.time.Duration; +import org.junit.After; +import org.junit.ClassRule; +import org.junit.Rule; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** End to end tests for {@link PAssert} and {@link Create} on Spark 4 Structured Streaming. */ +@RunWith(JUnit4.class) +@Category(StreamingTest.class) +public class PAssertStreamingTest implements Serializable { + + @ClassRule public static final SparkSessionRule SESSION = new SparkSessionRule(); + + @Rule public transient TemporaryFolder tempFolder = new TemporaryFolder(); + + @Rule public final transient TestPipeline pipeline = TestPipeline.create(); + + private static final String TAG_GLOBAL = "passert-global"; + private static final String TAG_UNEVEN = "passert-uneven"; + private static final String TAG_WINDOWED = "passert-windowed"; + private static final String TAG_EMPTY = "passert-empty"; + private static final String TAG_FAIL = "passert-fail"; + + private SparkStructuredStreamingPipelineOptions options() throws IOException { + SparkStructuredStreamingPipelineOptions options = + StreamingTestUtils.streamingOptions(tempFolder); + options.setTestMode(true); + return options; + } + + @After + public void tearDown() { + TestUnboundedSource.forget(TAG_GLOBAL); + TestUnboundedSource.forget(TAG_UNEVEN); + TestUnboundedSource.forget(TAG_WINDOWED); + TestUnboundedSource.forget(TAG_EMPTY); + TestUnboundedSource.forget(TAG_FAIL); + } + + @Test + public void testContainsInAnyOrderGlobalWindow() throws Exception { + PCollection out = + pipeline.apply("Read", Read.from(new TestUnboundedSource(TAG_GLOBAL, 2, 4))); + PAssert.that(out).containsInAnyOrder(TestUnboundedSource.elements(TAG_GLOBAL, 2, 4)); + pipeline.run(options()); + } + + @Test + public void testContainsInAnyOrderUnevenShards() throws Exception { + int[] shardCounts = {1, 4}; + PCollection out = + pipeline.apply("Read", Read.from(new TestUnboundedSource(TAG_UNEVEN, shardCounts))); + PAssert.that(out).containsInAnyOrder(TestUnboundedSource.elements(TAG_UNEVEN, shardCounts)); + SparkStructuredStreamingPipelineOptions options = options(); + options.setMaxRecordsPerBatch(2L); + pipeline.run(options); + } + + @Test + public void testContainsInAnyOrderFixedWindows() throws Exception { + PCollection out = + pipeline + .apply("Read", Read.from(new TestUnboundedSource(TAG_WINDOWED, 1, 4))) + .apply("Window", Window.into(FixedWindows.of(Duration.millis(2000)))); + PAssert.that(out).containsInAnyOrder(TestUnboundedSource.elements(TAG_WINDOWED, 1, 4)); + pipeline.run(options()); + } + + @Test + public void testEmpty() throws Exception { + PCollection out = + pipeline.apply("Read", Read.from(new TestUnboundedSource(TAG_EMPTY, 1, 0))); + PAssert.that(out).empty(); + pipeline.run(options()); + } + + @Test + public void testFailingAssertionFailsPipeline() throws Exception { + PCollection out = + pipeline.apply("Read", Read.from(new TestUnboundedSource(TAG_FAIL, 1, 2))); + PAssert.that(out).containsInAnyOrder("wrong"); + SparkStructuredStreamingPipelineOptions options = options(); + Pipeline.PipelineExecutionException thrown = + assertThrows(Pipeline.PipelineExecutionException.class, () -> pipeline.run(options)); + assertTrue( + "Expected PAssert AssertionError in causal chain of " + thrown, + Throwables.getCausalChain(thrown).stream() + .anyMatch( + t -> + t instanceof AssertionError + && t.getMessage() != null + && t.getMessage().contains("wrong"))); + } + + @Test + public void testCreateSingleElement() throws Exception { + PCollection out = pipeline.apply(Create.of("only")); + PAssert.that(out).containsInAnyOrder("only"); + pipeline.run(options()); + } + + @Test + public void testCreateMultipleElements() throws Exception { + PCollection out = pipeline.apply(Create.of("one", "two", "three")); + PAssert.that(out).containsInAnyOrder("one", "two", "three"); + pipeline.run(options()); + } +} diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTest.java index 5679950e1893..0e30d6ed791d 100644 --- a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTest.java +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTest.java @@ -125,11 +125,12 @@ public void testCheckpointRestartPreservesStateAndTimers() throws Exception { String c1 = StreamingTestUtils.newCollectorId("r1"); String c2 = StreamingTestUtils.newCollectorId("r2"); + TestUnboundedSource.extend(TAG_RESTART, 3); Pipeline p1 = pipeline( StreamingTestUtils.streamingOptions(checkpoint), TAG_RESTART, - 3, + 10, new RestartTestFn(), c1); assertEquals(PipelineResult.State.DONE, StreamingTestUtils.run(p1).getState()); diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineOptions.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineOptions.java index fb0192dba868..4d93fde3f8d1 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineOptions.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineOptions.java @@ -42,7 +42,8 @@ public interface SparkStructuredStreamingPipelineOptions extends SparkCommonPipe @Description( "Watermark delay in milliseconds applied to event timestamps of streaming sources " - + "(streaming mode only).") + + "(streaming mode only). Must be under one day for end of stream global window results " + + "to complete.") @Default.Long(0) long getWatermarkDelayMillis(); diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkTransformOverrides.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkTransformOverrides.java index a60aa59efd61..c53b89d4030a 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkTransformOverrides.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkTransformOverrides.java @@ -17,15 +17,34 @@ */ package org.apache.beam.runners.spark.structuredstreaming.translation; +import java.util.Collections; import java.util.List; +import java.util.Map; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.coders.ByteArrayCoder; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.io.BoundedSource; +import org.apache.beam.sdk.io.Read; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.runners.AppliedPTransform; import org.apache.beam.sdk.runners.PTransformOverride; +import org.apache.beam.sdk.runners.PTransformOverrideFactory; +import org.apache.beam.sdk.transforms.Impulse; import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.util.construction.PTransformMatchers; import org.apache.beam.sdk.util.construction.PTransformTranslation; +import org.apache.beam.sdk.util.construction.ReplacementOutputs; import org.apache.beam.sdk.util.construction.SplittableParDo; import org.apache.beam.sdk.util.construction.SplittableParDoNaiveBounded; +import org.apache.beam.sdk.util.construction.UnboundedReadFromBoundedSource.BoundedToUnboundedSourceAdapter; import org.apache.beam.sdk.util.construction.UnsupportedOverrideFactory; +import org.apache.beam.sdk.values.PBegin; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollection.IsBounded; +import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.joda.time.Instant; /** {@link PTransform} overrides for Spark runner. */ @SuppressWarnings({ @@ -50,7 +69,144 @@ public static List getDefaultOverrides(boolean streaming) { PTransformOverride.of( PTransformMatchers.urnEqualTo(PTransformTranslation.SPLITTABLE_PROCESS_KEYED_URN), new SplittableParDoNaiveBounded.OverrideFactory())); + } else { + builder + .add( + PTransformOverride.of( + PTransformMatchers.classEqualTo(SplittableParDo.PrimitiveBoundedRead.class), + new BoundedToUnboundedReadOverrideFactory())) + .add( + PTransformOverride.of( + PTransformMatchers.classEqualTo(Impulse.class), new ImpulseOverrideFactory())); } return builder.build(); } + + private static class BoundedToUnboundedReadOverrideFactory + implements PTransformOverrideFactory< + PBegin, PCollection, SplittableParDo.PrimitiveBoundedRead> { + @Override + public PTransformReplacement> getReplacementTransform( + AppliedPTransform, SplittableParDo.PrimitiveBoundedRead> + transform) { + return PTransformReplacement.of( + transform.getPipeline().begin(), + new StreamingBoundedRead<>(transform.getTransform().getSource())); + } + + @Override + public Map, ReplacementOutput> mapOutputs( + Map, PCollection> outputs, PCollection newOutput) { + return ReplacementOutputs.singleton(outputs, newOutput); + } + } + + private static class ImpulseOverrideFactory + implements PTransformOverrideFactory, Impulse> { + @Override + public PTransformReplacement> getReplacementTransform( + AppliedPTransform, Impulse> transform) { + return PTransformReplacement.of( + transform.getPipeline().begin(), new StreamingBoundedRead<>(new ImpulseSource())); + } + + @Override + public Map, ReplacementOutput> mapOutputs( + Map, PCollection> outputs, PCollection newOutput) { + return ReplacementOutputs.singleton(outputs, newOutput); + } + } + + /** + * Adapts a {@link BoundedSource} to an unbounded read while keeping the output bounded. {@link + * SplittableParDo.PrimitiveUnboundedRead} is applied directly because the {@link Read} to + * primitive conversion runs before {@code replaceTransforms}. + */ + private static class StreamingBoundedRead extends PTransform> { + private final BoundedSource source; + + StreamingBoundedRead(BoundedSource source) { + this.source = source; + source.validate(); + } + + @Override + public final PCollection expand(PBegin input) { + return Pipeline.applyTransform( + input, + new SplittableParDo.PrimitiveUnboundedRead<>( + Read.from(new BoundedToUnboundedSourceAdapter<>(source)))) + .setIsBoundedInternal(IsBounded.BOUNDED); + } + } + + /** Single element {@link BoundedSource} backing {@link Impulse} in streaming mode. */ + private static final class ImpulseSource extends BoundedSource { + private static final long serialVersionUID = 1L; + + @Override + public List> split( + long desiredBundleSizeBytes, PipelineOptions options) { + return Collections.singletonList(this); + } + + @Override + public long getEstimatedSizeBytes(PipelineOptions options) { + return 1L; + } + + @Override + public BoundedReader createReader(PipelineOptions options) { + return new ImpulseReader(this); + } + + @Override + public void validate() {} + + @Override + public Coder getOutputCoder() { + return ByteArrayCoder.of(); + } + } + + private static final class ImpulseReader extends BoundedSource.BoundedReader { + private final ImpulseSource source; + private boolean done; + + ImpulseReader(ImpulseSource source) { + this.source = source; + } + + @Override + public boolean start() { + if (done) { + return false; + } + done = true; + return true; + } + + @Override + public boolean advance() { + return false; + } + + @Override + public byte[] getCurrent() { + return new byte[0]; + } + + @Override + public Instant getCurrentTimestamp() { + return BoundedWindow.TIMESTAMP_MIN_VALUE; + } + + @Override + public void close() {} + + @Override + public BoundedSource getCurrentSource() { + return source; + } + } }