[SPARK-19876][SS][WIP] OneTime Trigger Executor

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/OffsetCommitLog.scala

@@ -0,0 +1,61 @@
+/*
+ * 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.spark.sql.execution.streaming
+
+import java.io.{InputStream, OutputStream}
+import java.nio.charset.StandardCharsets._
+
+import scala.io.{Source => IOSource}
+
+import org.apache.spark.sql.SparkSession
+
+class OffsetCommitLog(sparkSession: SparkSession, path: String)
+  extends HDFSMetadataLog[Option[String]](sparkSession, path) {
+
+  override protected def deserialize(in: InputStream): Option[String] = {
+    // called inside a try-finally where the underlying stream is closed in the caller
+    val lines = IOSource.fromInputStream(in, UTF_8.name()).getLines()
+    if (!lines.hasNext) {
+      throw new IllegalStateException("Incomplete log file")
+    }
+    val version = lines.next()
+    if (version != OffsetCommitLog.VERSION) {
+      throw new IllegalStateException(s"Unknown log version: ${version}")
+    }
+    // read metadata
+    lines.next().trim match {
+      case OffsetCommitLog.SERIALIZED_VOID => None
+      case metadata => Some(metadata)
+    }
+  }
+
+  override protected def serialize(metadata: Option[String], out: OutputStream): Unit = {
+    // called inside a try-finally where the underlying stream is closed in the caller
+    out.write(OffsetCommitLog.VERSION.getBytes(UTF_8))
+    out.write('\n')
+
+    // write metadata or void
+    out.write(metadata.getOrElse(OffsetCommitLog.SERIALIZED_VOID).getBytes(UTF_8))
+  }
+}
+
+object OffsetCommitLog {
+  private val VERSION = "v1"
+  private val SERIALIZED_VOID = "-"
+}
+

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StreamExecution.scala

@@ -162,6 +162,7 @@ class StreamExecution(
 
   private val triggerExecutor = trigger match {
     case t: ProcessingTime => ProcessingTimeExecutor(t, triggerClock)
+    case _ @ OneTime => OneTimeExecutor()
   }
 
   /** Defines the internal state of execution */
@@ -206,6 +207,12 @@ class StreamExecution(
    */
   val offsetLog = new OffsetSeqLog(sparkSession, checkpointFile("offsets"))
 
+  /**
+   * A log that records the committed batch ids. This is used to check if a batch was committed
+   * on restart, instead of (possibly) re-running the previous batch.
+   */
+  val commitLog = new OffsetCommitLog(sparkSession, checkpointFile("commits"))
+
   /** Whether all fields of the query have been initialized */
   private def isInitialized: Boolean = state.get != INITIALIZING
 
@@ -284,6 +291,7 @@ class StreamExecution(
               finishTrigger(dataAvailable)
               if (dataAvailable) {
                 // We'll increase currentBatchId after we complete processing current batch's data
+                commitLog.add(currentBatchId, None)
                 currentBatchId += 1
               } else {
                 currentStatus = currentStatus.copy(isDataAvailable = false)
@@ -377,17 +385,25 @@ class StreamExecution(
   private def populateStartOffsets(): Unit = {
     offsetLog.getLatest() match {
       case Some((batchId, nextOffsets)) =>
-        logInfo(s"Resuming streaming query, starting with batch $batchId")
-        currentBatchId = batchId
-        availableOffsets = nextOffsets.toStreamProgress(sources)
-        offsetSeqMetadata = nextOffsets.metadata.getOrElse(OffsetSeqMetadata())
-        logDebug(s"Found possibly unprocessed offsets $availableOffsets " +
-          s"at batch timestamp ${offsetSeqMetadata.batchTimestampMs}")
-
-        offsetLog.get(batchId - 1).foreach {
-          case lastOffsets =>
-            committedOffsets = lastOffsets.toStreamProgress(sources)
-            logDebug(s"Resuming with committed offsets: $committedOffsets")
+        currentBatchId = commitLog.getLatest() match {
+          case Some((committedBatchId, metadata)) => committedBatchId + 1
+          case None => batchId // can only happen if batchId is 0
+        }
+        if (currentBatchId > 0) {
+          offsetLog.get(currentBatchId - 1).foreach {
+            case lastOffsets =>
+              committedOffsets = lastOffsets.toStreamProgress(sources)
+              logDebug(s"Resuming with committed offsets: $committedOffsets")
+          }
+        }
+        logInfo(s"Resuming streaming query, starting with batch $currentBatchId")
+        if (currentBatchId == batchId) {
+          availableOffsets = nextOffsets.toStreamProgress(sources)
+          offsetSeqMetadata = nextOffsets.metadata.getOrElse(OffsetSeqMetadata())
+          logDebug(s"Found possibly unprocessed offsets $availableOffsets " +
+            s"at batch timestamp ${offsetSeqMetadata.batchTimestampMs}")
+        } else {
+          constructNextBatch()
         }
       case None => // We are starting this stream for the first time.
         logInfo(s"Starting new streaming query.")

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TriggerExecutor.scala

@@ -29,6 +29,17 @@ trait TriggerExecutor {
   def execute(batchRunner: () => Boolean): Unit
 }
 
+/**
+ * A trigger executor that runs a single batch only, then terminates.
+ */
+case class OneTimeExecutor() extends TriggerExecutor {
+
+  /**
+   * Execute a single batch using `batchRunner`.
+   */
+  override def execute(batchRunner: () => Boolean): Unit = batchRunner()
+}
+
 /**
  * A trigger executor that runs a batch every `intervalMs` milliseconds.
  */

sql/core/src/main/scala/org/apache/spark/sql/streaming/Trigger.scala

@@ -36,6 +36,26 @@ import org.apache.spark.unsafe.types.CalendarInterval
 @InterfaceStability.Evolving
 sealed trait Trigger
 
+/**
+ * :: Experimental ::
+ * A trigger that runs a query once then terminates
+ *
+ * Scala Example:
+ * {{{
+ *   df.write.trigger(OneTime)
+ * }}}
+ *
+ * Java Example:
+ * {{{
+ *   df.write.trigger(OneTime.create())
+ * }}}
+ *
+ * @since 2.2.0
+ */
+@Experimental
+@InterfaceStability.Evolving
+object OneTime extends Trigger
+
 /**
  * :: Experimental ::
  * A trigger that runs a query periodically based on the processing time. If `interval` is 0,

sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamSuite.scala

@@ -152,6 +152,15 @@ class StreamSuite extends StreamTest {
       AssertOnQuery(_.offsetLog.getLatest().get._1 == expectedId,
         s"offsetLog's latest should be $expectedId")
 
+    // Check the latest batchid in the commit log
+    def CheckCommitLogLatestBatchId(expectedId: Int): AssertOnQuery =
+      AssertOnQuery(_.commitLog.getLatest().get._1 == expectedId,
+        s"commitLog's latest should be $expectedId")
+
+    // Ensure that there has not been an incremental execution after restart
+    def CheckNoIncrementalExecutionCurrentBatchId(): AssertOnQuery =
+      AssertOnQuery(_.lastExecution == null, s"lastExecution not expected to run")
+
     // For each batch, we would log the state change during the execution
     // This checks whether the key of the state change log is the expected batch id
     def CheckIncrementalExecutionCurrentBatchId(expectedId: Int): AssertOnQuery =
@@ -177,6 +186,7 @@ class StreamSuite extends StreamTest {
       // Check the results of batch 0
       CheckAnswer(1, 2, 3),
       CheckIncrementalExecutionCurrentBatchId(0),
+      CheckCommitLogLatestBatchId(0),
       CheckOffsetLogLatestBatchId(0),
       CheckSinkLatestBatchId(0),
       // Add some data in batch 1
@@ -187,6 +197,7 @@ class StreamSuite extends StreamTest {
       // Check the results of batch 1
       CheckAnswer(1, 2, 3, 4, 5, 6),
       CheckIncrementalExecutionCurrentBatchId(1),
+      CheckCommitLogLatestBatchId(1),
       CheckOffsetLogLatestBatchId(1),
       CheckSinkLatestBatchId(1),
 
@@ -199,21 +210,23 @@ class StreamSuite extends StreamTest {
       // the currentId does not get logged (e.g. as 2) even if the clock has advanced many times
       CheckAnswer(1, 2, 3, 4, 5, 6),
       CheckIncrementalExecutionCurrentBatchId(1),
+      CheckCommitLogLatestBatchId(1),
       CheckOffsetLogLatestBatchId(1),
       CheckSinkLatestBatchId(1),
 
       /* Stop then restart the Stream  */
       StopStream,
       StartStream(ProcessingTime("10 seconds"), new StreamManualClock(60 * 1000)),
 
-      /* -- batch 1 rerun ----------------- */
-      // this batch 1 would re-run because the latest batch id logged in offset log is 1
+      /* -- batch 1 no rerun ----------------- */
+      // batch 1 would not re-run because the latest batch id logged in commit log is 1
       AdvanceManualClock(10 * 1000),
+      CheckNoIncrementalExecutionCurrentBatchId(),
 
       /* -- batch 2 ----------------------- */
       // Check the results of batch 1
       CheckAnswer(1, 2, 3, 4, 5, 6),
-      CheckIncrementalExecutionCurrentBatchId(1),
+      CheckCommitLogLatestBatchId(1),
       CheckOffsetLogLatestBatchId(1),
       CheckSinkLatestBatchId(1),
       // Add some data in batch 2
@@ -224,6 +237,7 @@ class StreamSuite extends StreamTest {
       // Check the results of batch 2
       CheckAnswer(1, 2, 3, 4, 5, 6, 7, 8, 9),
       CheckIncrementalExecutionCurrentBatchId(2),
+      CheckCommitLogLatestBatchId(2),
       CheckOffsetLogLatestBatchId(2),
       CheckSinkLatestBatchId(2))
   }

sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationSuite.scala

@@ -277,7 +277,8 @@ class StreamingAggregationSuite extends StateStoreMetricsTest with BeforeAndAfte
         true
       },
       StartStream(ProcessingTime("10 seconds"), triggerClock = clock),
-      CheckLastBatch((20L, 1), (85L, 1)),
+      // The commit log should ensure that we do not run another batch
+      CheckLastBatch(),
       AssertOnQuery { q =>
         clock.getTimeMillis() == 90000L
       },
@@ -327,7 +328,8 @@ class StreamingAggregationSuite extends StateStoreMetricsTest with BeforeAndAfte
         true
       },
       StartStream(ProcessingTime("10 day"), triggerClock = clock),
-      CheckLastBatch((20L, 1), (85L, 1)),
+      // Commit log should prevent batch from running again
+      CheckLastBatch(),
 
       // advance clock to 100 days, should retain keys >= 90
       AddData(inputData, 85L, 90L, 100L, 105L),

sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryListenerSuite.scala

@@ -109,8 +109,9 @@ class StreamingQueryListenerSuite extends StreamTest with BeforeAndAfter {
 
         // Termination event generated with exception message when stopped with error
         StartStream(ProcessingTime(100), triggerClock = clock),
+        AdvanceManualClock(100), // advance clock to ensure completed initial trigger
         AddData(inputData, 0),
-        AdvanceManualClock(100),
+        AdvanceManualClock(100), // process bad data
         ExpectFailure[SparkException](),
         AssertOnQuery { query =>
           eventually(Timeout(streamingTimeout)) {

sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQuerySuite.scala

@@ -158,6 +158,51 @@ class StreamingQuerySuite extends StreamTest with BeforeAndAfter with Logging wi
     )
   }
 
+  testQuietly("OneTime trigger, commit log, and exception") {
+    val inputData = MemoryStream[Int]
+    val mapped = inputData.toDS().map { 6 / _}
+
+    testStream(mapped)(
+      AssertOnQuery(_.isActive === true),
+      StopStream,
+      AddData(inputData, 1, 2),
+      StartStream(trigger = OneTime),
+      CheckAnswer(6, 3),
+      AssertOnQuery(_.isActive === false),
+      StopStream, // clears out StreamTest state
+      AssertOnQuery { q =>
+        // both commit log and offset log contain the same (latest) batch id
+        q.commitLog.getLatest().map(_._1).getOrElse(-1L) ==
+          q.offsetLog.getLatest().map(_._1).getOrElse(-2L)
+      },
+      AssertOnQuery { q =>
+        // blow away commit log and sink result
+        q.commitLog.purge(1)
+        q.sink.asInstanceOf[MemorySink].clear()
+        true
+      },
+      StartStream(trigger = OneTime),
+      CheckAnswer(6, 3), // ensure we fall back to offset log and reprocess batch
+      AssertOnQuery(_.isActive === false),
+      StopStream,
+      AddData(inputData, 3),
+      StartStream(trigger = OneTime),
+      CheckLastBatch(2), // commit log should be back in place
+      AssertOnQuery(_.isActive === false),
+      StopStream,
+      AddData(inputData, 0),
+      StartStream(OneTime),
+      ExpectFailure[SparkException](),
+      AssertOnQuery(_.isActive === false),
+      AssertOnQuery(q => {
+        q.exception.get.startOffset ===
+          q.committedOffsets.toOffsetSeq(Seq(inputData), OffsetSeqMetadata()).toString &&
+          q.exception.get.endOffset ===
+            q.availableOffsets.toOffsetSeq(Seq(inputData), OffsetSeqMetadata()).toString
+      }, "incorrect start offset or end offset on exception")
+    )
+  }
+
   testQuietly("status, lastProgress, and recentProgress") {
     import StreamingQuerySuite._
     clock = new StreamManualClock
@@ -306,8 +351,9 @@ class StreamingQuerySuite extends StreamTest with BeforeAndAfter with Logging wi
 
       // Test status and progress after query terminated with error
       StartStream(ProcessingTime(100), triggerClock = clock),
+      AdvanceManualClock(100), // ensure initial trigger completes before AddData
       AddData(inputData, 0),
-      AdvanceManualClock(100),
+      AdvanceManualClock(100), // allow another trigger
       ExpectFailure[SparkException](),
       AssertOnQuery(_.status.isDataAvailable === false),
       AssertOnQuery(_.status.isTriggerActive === false),