This is an automated email from the ASF dual-hosted git repository.
Abacn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new a757f49cb50 [#36841] Spark 4 streaming: Impulse, bounded reads and
PAssert on the Structured Streaming runner (#40766)
a757f49cb50 is described below
commit a757f49cb50631f0ab7341487eab11cedd3811e2
Author: Tobias Kaymak <[email protected]>
AuthorDate: Wed Oct 28 16:21:00 2026 +0100
[#36841] Spark 4 streaming: Impulse, bounded reads and PAssert on the
Structured Streaming runner (#40766)
* [#36841] Spark 4 streaming: Impulse, bounded reads and PAssert on the
Structured Streaming runner
* Run the Spark ValidatesRunner PostCommits
---
...eam_PostCommit_Java_ValidatesRunner_Spark4.json | 2 +-
...a_ValidatesRunner_SparkStructuredStreaming.json | 2 +-
.../io/streaming/BeamReaderCache.java | 37 ++++-
.../io/streaming/BeamSourceCheckpoint.java | 32 ++++-
.../io/streaming/UnboundedSourceDataset.java | 35 ++++-
.../translation/PipelineTranslatorStreaming.java | 16 +--
.../translation/StreamingEvaluationContext.java | 3 +-
.../streaming/ReadUnboundedTranslator.java | 20 ++-
.../io/streaming/BeamMicroBatchSourceTest.java | 18 ++-
.../io/streaming/TestUnboundedSource.java | 37 ++++-
.../PipelineTranslatorStreamingTest.java | 17 ---
.../streaming/PAssertStreamingTest.java | 149 ++++++++++++++++++++
.../streaming/StatefulParDoStreamingTest.java | 3 +-
.../SparkStructuredStreamingPipelineOptions.java | 3 +-
.../translation/SparkTransformOverrides.java | 156 +++++++++++++++++++++
15 files changed, 470 insertions(+), 60 deletions(-)
diff --git
a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark4.json
b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark4.json
index e0266d62f2e..f1ba03a243e 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 1b832d7f19c..65728e1581d 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 e961ff2f380..8348d04d289 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;
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 final class BeamReaderCache {
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 final class BeamReaderCache {
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 @@ public final class BeamReaderCache {
if (positionEpoch != epoch || moved) {
return false;
}
+ doneEligible |= positionExhausted;
inBatch = true;
lastUsedMillis = System.currentTimeMillis();
return true;
@@ -298,6 +325,8 @@ public final class BeamReaderCache {
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 d667976edce..8b58a0e5107 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 @@ import org.slf4j.LoggerFactory;
*
* <p>{@code <location>/splits} pins the split list, written once by the
driver. {@code
* <location>/marks/<epoch>/<splitId>} 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 <location>/done/<splitId>} 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 final class BeamSourceCheckpoint {
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 final class BeamSourceCheckpoint {
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 <root>/sources/<index>}, 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 33106824c10..3ef35762748 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.coders.Coder;
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 @@ import scala.reflect.ClassTag;
* TIMESTAMP} holding the event timestamp of that element.
*
* <p>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 @@ public final class UnboundedSourceDataset {
private final AtomicLong purgeRequested = new AtomicLong();
private long epoch;
+ private boolean emitSentinel;
private @Nullable List<UnboundedSource<T, ?>> splits;
BeamMicroBatchStream(BeamTable<T> table, String checkpointLocation) {
@@ -343,6 +347,13 @@ public final class UnboundedSourceDataset {
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 final class UnboundedSourceDataset {
endEpoch,
quotas[i],
table.maxBatchDurationMillis,
- table.readerIdleTimeoutMillis);
+ table.readerIdleTimeoutMillis,
+ emitSentinel);
}
return partitions;
}
@@ -495,6 +507,7 @@ public final class UnboundedSourceDataset {
final long maxBatchDurationMillis;
final long readerIdleTimeoutMillis;
+ final boolean emitSentinel;
BeamInputPartition(
UnboundedSource<T, ?> split,
@@ -507,7 +520,8 @@ public final class UnboundedSourceDataset {
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 @@ public final class UnboundedSourceDataset {
this.maxRecords = maxRecords;
this.maxBatchDurationMillis = maxBatchDurationMillis;
this.readerIdleTimeoutMillis = readerIdleTimeoutMillis;
+ this.emitSentinel = emitSentinel;
}
@Override
@@ -573,6 +588,7 @@ public final class UnboundedSourceDataset {
private final long endEpoch;
private final long maxRecords;
private final long maxBatchDurationMillis;
+ private final boolean emitSentinel;
private long recordsRead;
private long deadlineMillis = -1L;
@@ -586,6 +602,7 @@ public final class UnboundedSourceDataset {
this.endEpoch = partition.endEpoch;
this.maxRecords = partition.maxRecords;
this.maxBatchDurationMillis = partition.maxBatchDurationMillis;
+ this.emitSentinel = partition.emitSentinel;
PipelineOptions options = partition.options.value().get();
Configuration conf = partition.hadoopConf.value().value();
BeamSourceCheckpoint checkpoint =
@@ -624,6 +641,10 @@ public final class UnboundedSourceDataset {
current = toRow();
return true;
}
+ if (cached.claimSentinel(emitSentinel)) {
+ current = sentinelRow();
+ return true;
+ }
if (backOff == null) {
backOff = backOff(remaining);
}
@@ -684,6 +705,9 @@ public final class UnboundedSourceDataset {
CheckpointMark mark = cached.reader().getCheckpointMark();
byte[] codedMark = encodeMark(split, mark);
checkpoint.writeMark(splitId, endEpoch, codedMark);
+ if (cached.shouldWriteDone()) {
+ checkpoint.writeDone(splitId);
+ }
cached.endBatch(endEpoch, mark, codedMark);
LOG.debug("Beam reader {} read {} record(s) up to epoch {}.", key,
recordsRead, endEpoch);
}
@@ -717,5 +741,10 @@ public final class UnboundedSourceDataset {
// Spark stores TimestampType as microseconds.
return new GenericInternalRow(new Object[] {payload,
timestamp.getMillis() * 1000L});
}
+
+ private static InternalRow sentinelRow() {
+ return new GenericInternalRow(
+ new Object[] {new byte[0],
BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis() * 1000L});
+ }
}
}
diff --git
a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java
b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java
index 6f2d14a69b8..4d93741b636 100644
---
a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java
+++
b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java
@@ -29,12 +29,12 @@ import org.apache.beam.sdk.state.TimerSpec;
import org.apache.beam.sdk.transforms.Combine;
import org.apache.beam.sdk.transforms.DoFn;
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.transforms.windowing.DefaultTrigger;
+import org.apache.beam.sdk.transforms.windowing.Never.NeverTrigger;
import org.apache.beam.sdk.util.construction.SplittableParDo;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
@@ -93,17 +93,6 @@ public class PipelineTranslatorStreaming extends
PipelineTranslatorCommon {
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 @@ public class PipelineTranslatorStreaming extends
PipelineTranslatorCommon {
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 2ca54a1e948..b5057cd260a 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.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.Pipeline;
import org.apache.beam.sdk.annotations.Internal;
import org.apache.hadoop.fs.Path;
import org.apache.spark.sql.Dataset;
@@ -179,7 +180,7 @@ public class StreamingEvaluationContext extends
EvaluationContext {
} 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 eb1421b4554..2281bbfdec3 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 @@ package
org.apache.beam.runners.spark.structuredstreaming.translation.streaming;
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.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.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 @@ public class ReadUnboundedTranslator<T>
Dataset<WindowedValue<T>> 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<T> implements MapFunction<byte[],
WindowedValue<T>> {
+ /**
+ * Decodes the binary payload column back into a Beam {@code WindowedValue},
dropping the empty
+ * end of stream sentinel payload.
+ */
+ private static final class DecodePayload<T> implements
FlatMapFunction<byte[], WindowedValue<T>> {
private final Coder<WindowedValue<T>> coder;
DecodePayload(Coder<WindowedValue<T>> coder) {
@@ -90,8 +95,11 @@ public class ReadUnboundedTranslator<T>
}
@Override
- public WindowedValue<T> call(byte[] payload) {
- return CoderHelpers.fromByteArray(payload, coder);
+ public Iterator<WindowedValue<T>> 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 559deb80803..07c823068c7 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.Dataset;
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 @@ public class BeamMicroBatchSourceTest implements
Serializable {
(batch, batchId) -> {
List<List<Row>> target = BATCHES.get(queryName);
if (target != null) {
- target.add(batch.collectAsList());
+ List<Row> 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 @@ public class BeamMicroBatchSourceTest implements
Serializable {
end,
3L,
30_000L,
- 600_000L);
+ 600_000L,
+ false);
}
private static List<String> readBatch(BeamInputPartition<String> partition)
throws IOException {
@@ -778,8 +784,10 @@ public class BeamMicroBatchSourceTest implements
Serializable {
private static void drainInto(BeamPartitionReader<String> reader,
List<String> 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 98f62501f37..8bc257d1872 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 @@ 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.
+ * 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<String,
TestUnboundedSource.Mark> {
private static final long serialVersionUID = 1L;
@@ -62,20 +63,27 @@ public final class TestUnboundedSource extends
UnboundedSource<String, TestUnbou
private final int shard;
private final int shards;
private final int perShard;
+ private final int @Nullable [] shardCounts;
public TestUnboundedSource(String tag, int shards, int count) {
- this(tag, -1, shards, count / shards);
+ this(tag, -1, shards, count / shards, null);
+ }
+
+ public TestUnboundedSource(String tag, int[] shardCounts) {
+ this(tag, -1, shardCounts.length, 0, shardCounts.clone());
}
public static void extend(String tag, int count) {
EXTENDED.put(tag, count);
}
- private TestUnboundedSource(String tag, int shard, int shards, int perShard)
{
+ private TestUnboundedSource(
+ String tag, int shard, int shards, int perShard, int @Nullable []
shardCounts) {
this.tag = tag;
this.shard = shard;
this.shards = shards;
this.perShard = perShard;
+ this.shardCounts = shardCounts;
}
public static Set<String> elements(String tag, int shards, int count) {
@@ -88,6 +96,16 @@ public final class TestUnboundedSource extends
UnboundedSource<String, TestUnbou
return elements;
}
+ public static Set<String> elements(String tag, int[] shardCounts) {
+ Set<String> 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 final class TestUnboundedSource extends
UnboundedSource<String, TestUnbou
}
List<TestUnboundedSource> 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 @@ public final class TestUnboundedSource extends
UnboundedSource<String, TestUnbou
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 final class TestUnboundedSource extends
UnboundedSource<String, TestUnbou
current = next++;
return true;
}
+ exhausted = next >= source.perShard;
return false;
}
@@ -238,11 +259,15 @@ public final class TestUnboundedSource extends
UnboundedSource<String, TestUnbou
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 a5afa7c1a4a..d18f069cfbe 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.TimeDomain;
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 @@ public class PipelineTranslatorStreamingTest implements
Serializable {
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 00000000000..79a836f1c02
--- /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<String> 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<String> 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<String> 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<String> 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<String> 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<String> out = pipeline.apply(Create.of("only"));
+ PAssert.that(out).containsInAnyOrder("only");
+ pipeline.run(options());
+ }
+
+ @Test
+ public void testCreateMultipleElements() throws Exception {
+ PCollection<String> 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 5679950e189..0e30d6ed791 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 class StatefulParDoStreamingTest implements
Serializable {
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 fb0192dba86..4d93fde3f8d 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 a60aa59efd6..c53b89d4030 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 @@ class SparkTransformOverrides {
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<T>
+ implements PTransformOverrideFactory<
+ PBegin, PCollection<T>, SplittableParDo.PrimitiveBoundedRead<T>> {
+ @Override
+ public PTransformReplacement<PBegin, PCollection<T>>
getReplacementTransform(
+ AppliedPTransform<PBegin, PCollection<T>,
SplittableParDo.PrimitiveBoundedRead<T>>
+ transform) {
+ return PTransformReplacement.of(
+ transform.getPipeline().begin(),
+ new StreamingBoundedRead<>(transform.getTransform().getSource()));
+ }
+
+ @Override
+ public Map<PCollection<?>, ReplacementOutput> mapOutputs(
+ Map<TupleTag<?>, PCollection<?>> outputs, PCollection<T> newOutput) {
+ return ReplacementOutputs.singleton(outputs, newOutput);
+ }
+ }
+
+ private static class ImpulseOverrideFactory
+ implements PTransformOverrideFactory<PBegin, PCollection<byte[]>,
Impulse> {
+ @Override
+ public PTransformReplacement<PBegin, PCollection<byte[]>>
getReplacementTransform(
+ AppliedPTransform<PBegin, PCollection<byte[]>, Impulse> transform) {
+ return PTransformReplacement.of(
+ transform.getPipeline().begin(), new StreamingBoundedRead<>(new
ImpulseSource()));
+ }
+
+ @Override
+ public Map<PCollection<?>, ReplacementOutput> mapOutputs(
+ Map<TupleTag<?>, PCollection<?>> outputs, PCollection<byte[]>
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<T> extends PTransform<PBegin,
PCollection<T>> {
+ private final BoundedSource<T> source;
+
+ StreamingBoundedRead(BoundedSource<T> source) {
+ this.source = source;
+ source.validate();
+ }
+
+ @Override
+ public final PCollection<T> 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<byte[]> {
+ private static final long serialVersionUID = 1L;
+
+ @Override
+ public List<? extends BoundedSource<byte[]>> split(
+ long desiredBundleSizeBytes, PipelineOptions options) {
+ return Collections.singletonList(this);
+ }
+
+ @Override
+ public long getEstimatedSizeBytes(PipelineOptions options) {
+ return 1L;
+ }
+
+ @Override
+ public BoundedReader<byte[]> createReader(PipelineOptions options) {
+ return new ImpulseReader(this);
+ }
+
+ @Override
+ public void validate() {}
+
+ @Override
+ public Coder<byte[]> getOutputCoder() {
+ return ByteArrayCoder.of();
+ }
+ }
+
+ private static final class ImpulseReader extends
BoundedSource.BoundedReader<byte[]> {
+ 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<byte[]> getCurrentSource() {
+ return source;
+ }
+ }
}
