[#36841] Spark 4 streaming: Impulse, bounded reads and PAssert on the Structured Streaming runner

.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
 }

.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"
 }

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<T> 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<T> 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();

runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamSourceCheckpoint.java

@@ -42,7 +42,8 @@
  *
  * <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 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 <root>/sources/<index>}, the batch id is the

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.
  *
  * <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 @@ static class BeamMicroBatchStream<T> implements MicroBatchStream {
     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 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<T> implements InputPartition {
 
     final long maxBatchDurationMillis;
     final long readerIdleTimeoutMillis;
+    final boolean emitSentinel;
 
     BeamInputPartition(
         UnboundedSource<T, ?> split,
@@ -507,7 +520,8 @@ static final class BeamInputPartition<T> 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<T> implements InputPartition {
       this.maxRecords = maxRecords;
       this.maxBatchDurationMillis = maxBatchDurationMillis;
       this.readerIdleTimeoutMillis = readerIdleTimeoutMillis;
+      this.emitSentinel = emitSentinel;
     }
 
     @Override
@@ -573,6 +588,7 @@ static final class BeamPartitionReader<T> implements PartitionReader<InternalRow
     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 @@ static final class BeamPartitionReader<T> implements PartitionReader<InternalRow
       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 boolean next() throws IOException {
           current = toRow();
           return true;
         }
+        if (cached.claimSentinel(emitSentinel)) {
+          current = sentinelRow();
+          return true;
+        }
         if (backOff == null) {
           backOff = backOff(remaining);
         }
@@ -684,6 +705,9 @@ private void endBatch(boolean discarded) throws IOException {
       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 @@ private InternalRow toRow() {
       // 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});
+    }
   }
 }

runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java

@@ -29,12 +29,12 @@
 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 @@ TransformTranslator<InT, OutT, TransformT> 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);
     }

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<StreamingQuery> 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);
         }
       }
     }

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,22 +78,28 @@ protected void translate(SplittableParDo.PrimitiveUnboundedRead<T> transform, Co
     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) {
       this.coder = coder;
     }
 
     @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();
     }
   }
 }