[SPARK-14176][SQL]Add DataFrameWriter.trigger to set the stream batch period

sql/core/src/main/scala/org/apache/spark/sql/ContinuousQueryManager.scala

@@ -171,13 +171,20 @@ class ContinuousQueryManager(sqlContext: SQLContext) {
       name: String,
       checkpointLocation: String,
       df: DataFrame,
-      sink: Sink): ContinuousQuery = {
+      sink: Sink,
+      triggerIntervalMs: Long): ContinuousQuery = {
     activeQueriesLock.synchronized {
       if (activeQueries.contains(name)) {
         throw new IllegalArgumentException(
           s"Cannot start query with name $name as a query with that name is already active")
       }
-      val query = new StreamExecution(sqlContext, name, checkpointLocation, df.logicalPlan, sink)
+      val query = new StreamExecution(
+        sqlContext,
+        name,
+        checkpointLocation,
+        df.logicalPlan,
+        sink,
+        triggerIntervalMs)
       query.start()
       activeQueries.put(name, query)
       query

sql/core/src/main/scala/org/apache/spark/sql/DataFrameWriter.scala

@@ -20,6 +20,7 @@ package org.apache.spark.sql
 import java.util.Properties
 
 import scala.collection.JavaConverters._
+import scala.concurrent.duration._
 
 import org.apache.hadoop.fs.Path
 
@@ -77,6 +78,26 @@ final class DataFrameWriter private[sql](df: DataFrame) {
     this
   }
 
+  /**
+   * Set the trigger interval for the stream query.
+   *
+   * @since 2.0.0
+   */
+  def trigger(interval: Duration): DataFrameWriter = {
+    this.extraOptions += ("triggerInterval" -> interval.toMillis.toString)
+    this
+  }
+
+  /**
+   * Set the trigger interval for the stream query.
+   *
+   * @since 2.0.0
+   */
+  def trigger(interval: Long, unit: TimeUnit): DataFrameWriter = {
+    this.extraOptions += ("triggerInterval" -> unit.toMillis(interval).toString)
+    this
+  }
+
   /**
    * Specifies the underlying output data source. Built-in options include "parquet", "json", etc.
    *
@@ -257,11 +278,14 @@ final class DataFrameWriter private[sql](df: DataFrame) {
     val checkpointLocation = extraOptions.getOrElse("checkpointLocation", {
       new Path(df.sqlContext.conf.checkpointLocation, queryName).toUri.toString
     })
+    val triggerIntervalMs = extraOptions.getOrElse("triggerInterval", "0").toLong
+    require(triggerIntervalMs >= 0, "the interval of trigger should not be negative")
     df.sqlContext.sessionState.continuousQueryManager.startQuery(
       queryName,
       checkpointLocation,
       df,
-      dataSource.createSink())
+      dataSource.createSink(),
+      triggerIntervalMs)
   }
 
   /**

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

@@ -46,16 +46,14 @@ class StreamExecution(
     override val name: String,
     val checkpointRoot: String,
     private[sql] val logicalPlan: LogicalPlan,
-    val sink: Sink) extends ContinuousQuery with Logging {
+    val sink: Sink,
+    triggerIntervalMs: Long) extends ContinuousQuery with Logging {
 
   /** An monitor used to wait/notify when batches complete. */
   private val awaitBatchLock = new Object
   private val startLatch = new CountDownLatch(1)
   private val terminationLatch = new CountDownLatch(1)
 
-  /** Minimum amount of time in between the start of each batch. */
-  private val minBatchTime = 10
-
   /**
    * Tracks how much data we have processed and committed to the sink or state store from each
    * input source.
@@ -212,9 +210,18 @@ class StreamExecution(
       populateStartOffsets()
       logDebug(s"Stream running from $committedOffsets to $availableOffsets")
       while (isActive) {
+        val batchStartTimeMs = System.currentTimeMillis()
         if (dataAvailable) runBatch()
         commitAndConstructNextBatch()
-        Thread.sleep(minBatchTime) // TODO: Could be tighter
+        if (triggerIntervalMs > 0) {
+          val batchElapsedTime = System.currentTimeMillis() - batchStartTimeMs
+          if (batchElapsedTime > triggerIntervalMs) {
+            logWarning("Current batch is falling behind. The trigger interval is " +
+              s"${triggerIntervalMs} milliseconds, but spent ${batchElapsedTime} milliseconds")
+          } else {
+            Thread.sleep(triggerIntervalMs - batchElapsedTime)
+          }
+        }
       }
     } catch {
       case _: InterruptedException if state == TERMINATED => // interrupted by stop()

sql/core/src/test/scala/org/apache/spark/sql/StreamTest.scala

@@ -276,7 +276,7 @@ trait StreamTest extends QueryTest with Timeouts {
             currentStream =
               sqlContext
                 .streams
-                .startQuery(StreamExecution.nextName, metadataRoot, stream, sink)
+                .startQuery(StreamExecution.nextName, metadataRoot, stream, sink, 10L)
                 .asInstanceOf[StreamExecution]
             currentStream.microBatchThread.setUncaughtExceptionHandler(
               new UncaughtExceptionHandler {

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

@@ -243,7 +243,8 @@ class ContinuousQueryManagerSuite extends StreamTest with SharedSQLContext with
                 StreamExecution.nextName,
                 metadataRoot,
                 df,
-                new MemorySink(df.schema))
+                new MemorySink(df.schema),
+                10L)
               .asInstanceOf[StreamExecution]
           } catch {
             case NonFatal(e) =>

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

@@ -17,6 +17,10 @@
 
 package org.apache.spark.sql.streaming.test
 
+import java.util.concurrent.TimeUnit
+
+import scala.concurrent.duration._
+
 import org.scalatest.BeforeAndAfter
 
 import org.apache.spark.sql._
@@ -274,4 +278,28 @@ class DataFrameReaderWriterSuite extends StreamTest with SharedSQLContext with B
     assert(activeStreamNames.contains("name"))
     sqlContext.streams.active.foreach(_.stop())
   }
+
+  test("trigger") {
+    val df = sqlContext.read
+      .format("org.apache.spark.sql.streaming.test")
+      .stream("/test")
+
+    df.write
+      .format("org.apache.spark.sql.streaming.test")
+      .option("checkpointLocation", newMetadataDir)
+      .trigger(10.seconds)
+      .startStream()
+      .stop()
+
+    assert(LastOptions.parameters("triggerInterval") == "10000")
+
+    df.write
+      .format("org.apache.spark.sql.streaming.test")
+      .option("checkpointLocation", newMetadataDir)
+      .trigger(100, TimeUnit.SECONDS)
+      .startStream()
+      .stop()
+
+    assert(LastOptions.parameters("triggerInterval") == "100000")
+  }
 }