Address and fix timestamp for rampUpTimeSeconds

Author: zsxwing

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

@@ -24,6 +24,7 @@ import java.util.concurrent.TimeUnit
 import org.apache.commons.io.IOUtils
 
 import org.apache.spark.internal.Logging
+import org.apache.spark.network.util.JavaUtils
 import org.apache.spark.sql.{DataFrame, SQLContext}
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.util.{CaseInsensitiveMap, DateTimeUtils}
@@ -37,11 +38,12 @@ import org.apache.spark.util.{ManualClock, SystemClock}
  *  with 0L.
  *
  *  This source supports the following options:
- *  - `tuplesPerSecond` (default: 1): How many tuples should be generated per second.
- *  - `rampUpTimeSeconds` (default: 0): How many seconds to ramp up before the generating speed
- *    becomes `tuplesPerSecond`.
- *  - `numPartitions` (default: Spark's default parallelism): The partition number for the generated
- *    tuples.
+ *  - `tuplesPerSecond` (e.g. 100, default: 1): How many tuples should be generated per second.
+ *  - `rampUpTime` (e.g. 5s, default: 0s): How long to ramp up before the generating speed
+ *    becomes `tuplesPerSecond`. Using finer granularities than seconds will be truncated to integer
+ *    seconds.
+ *  - `numPartitions` (e.g. 10, default: Spark's default parallelism): The partition number for the
+ *    generated tuples.
  */
 class RateSourceProvider extends StreamSourceProvider with DataSourceRegister {
 
@@ -63,22 +65,23 @@ class RateSourceProvider extends StreamSourceProvider with DataSourceRegister {
     val tuplesPerSecond = params.get("tuplesPerSecond").map(_.toLong).getOrElse(1L)
     if (tuplesPerSecond <= 0) {
       throw new IllegalArgumentException(
-        s"Invalid value '${params("tuplesPerSecond")}' for option 'tuplesPerSecond', " +
+        s"Invalid value '${params("tuplesPerSecond")}'. The option 'tuplesPerSecond' " +
           "must be positive")
     }
 
-    val rampUpTimeSeconds = params.get("rampUpTimeSeconds").map(_.toLong).getOrElse(0L)
+    val rampUpTimeSeconds =
+      params.get("rampUpTime").map(JavaUtils.timeStringAsSec(_)).getOrElse(0L)
     if (rampUpTimeSeconds < 0) {
       throw new IllegalArgumentException(
-        s"Invalid value '${params("rampUpTimeSeconds")}' for option 'rampUpTimeSeconds', " +
+        s"Invalid value '${params("rampUpTime")}'. The option 'rampUpTime' " +
           "must not be negative")
     }
 
     val numPartitions = params.get("numPartitions").map(_.toInt).getOrElse(
       sqlContext.sparkContext.defaultParallelism)
     if (numPartitions <= 0) {
       throw new IllegalArgumentException(
-        s"Invalid value '${params("numPartitions")}' for option 'numPartitions', " +
+        s"Invalid value '${params("numPartitions")}'. The option 'numPartitions' " +
           "must be positive")
     }
 
@@ -117,8 +120,9 @@ class RateStreamSource(
   private val maxSeconds = Long.MaxValue / tuplesPerSecond
 
   if (rampUpTimeSeconds > maxSeconds) {
-    throw new ArithmeticException("integer overflow. Max offset with tuplesPerSecond " +
-      s"$tuplesPerSecond is $maxSeconds, but 'rampUpTimeSeconds' is $rampUpTimeSeconds.")
+    throw new ArithmeticException(
+      s"Integer overflow. Max offset with $tuplesPerSecond tuplesPerSecond" +
+        s" is $maxSeconds, but 'rampUpTimeSeconds' is $rampUpTimeSeconds.")
   }
 
   private val startTimeMs = {
@@ -175,10 +179,10 @@ class RateStreamSource(
   override def getBatch(start: Option[Offset], end: Offset): DataFrame = {
     val startSeconds = start.flatMap(LongOffset.convert(_).map(_.offset)).getOrElse(0L)
     val endSeconds = LongOffset.convert(end).map(_.offset).getOrElse(0L)
-    assert(startSeconds <= endSeconds)
+    assert(startSeconds <= endSeconds, s"startSeconds($startSeconds) > endSeconds($endSeconds)")
     if (endSeconds > maxSeconds) {
-      throw new ArithmeticException("integer overflow. Max offset with " +
-        s"tuplesPerSecond $tuplesPerSecond is $maxSeconds, but it's $endSeconds now.")
+      throw new ArithmeticException("Integer overflow. Max offset with " +
+        s"$tuplesPerSecond tuplesPerSecond is $maxSeconds, but it's $endSeconds now.")
     }
     // Fix "lastTimeMs" for recovery
     if (lastTimeMs < TimeUnit.SECONDS.toMillis(endSeconds) + startTimeMs) {
@@ -188,11 +192,17 @@ class RateStreamSource(
     val rangeEnd = valueAtSecond(endSeconds, tuplesPerSecond, rampUpTimeSeconds)
     logDebug(s"startSeconds: $startSeconds, endSeconds: $endSeconds, " +
       s"rangeStart: $rangeStart, rangeEnd: $rangeEnd")
-    val localStartTimeMs = startTimeMs
-    val localPerSecond = tuplesPerSecond
+
+    if (rangeStart == rangeEnd) {
+      return sqlContext.internalCreateDataFrame(sqlContext.sparkContext.emptyRDD, schema)
+    }
+
+    val localStartTimeMs = startTimeMs + TimeUnit.SECONDS.toMillis(startSeconds)
+    val relativeMsPerValue =
+      TimeUnit.SECONDS.toMillis(endSeconds - startSeconds) / (rangeEnd - rangeStart)
 
     val rdd = sqlContext.sparkContext.range(rangeStart, rangeEnd, 1, numPartitions).map { v =>
-      val relative = v * 1000L / localPerSecond
+      val relative = (v - rangeStart) * relativeMsPerValue
       InternalRow(DateTimeUtils.fromMillis(relative + localStartTimeMs), v)
     }
     sqlContext.internalCreateDataFrame(rdd, schema)

sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/RateSourceSuite.scala

@@ -39,14 +39,6 @@ class RateSourceSuite extends StreamTest {
     }
   }
 
-  private def getManualClockFromQuery(query: StreamExecution): ManualClock = {
-    val rateSource = query.logicalPlan.collect {
-      case StreamingExecutionRelation(source, _) if source.isInstanceOf[RateStreamSource] =>
-        source.asInstanceOf[RateStreamSource]
-    }.head
-    rateSource.clock.asInstanceOf[ManualClock]
-  }
-
   test("basic") {
     val input = spark.readStream
       .format("rate")
@@ -80,30 +72,31 @@ class RateSourceSuite extends StreamTest {
     assert(valueAtSecond(seconds = 5, tuplesPerSecond = 10, rampUpTimeSeconds = 4) === 30)
   }
 
-  test("rampUpTimeSeconds") {
+  test("rampUpTime") {
     val input = spark.readStream
       .format("rate")
       .option("tuplesPerSecond", "10")
-      .option("rampUpTimeSeconds", "4")
+      .option("rampUpTime", "4s")
       .option("useManualClock", "true")
       .load()
-      .select($"value")
+      .as[(java.sql.Timestamp, Long)]
+      .map(v => (v._1.getTime, v._2))
     testStream(input)(
       AdvanceRateManualClock(seconds = 1),
-      CheckLastBatch(0 until 2: _*), // speed = 2
+      CheckLastBatch((0 until 2).map(v => v * 500 -> v): _*), // speed = 2
       AdvanceRateManualClock(seconds = 1),
-      CheckLastBatch(2 until 6: _*), // speed = 4
+      CheckLastBatch((2 until 6).map(v => 1000 + (v - 2) * 250 -> v): _*), // speed = 4
       AdvanceRateManualClock(seconds = 1),
-      CheckLastBatch(6 until 12: _*), // speed = 6
+      CheckLastBatch((6 until 12).map(v => 2000 + (v - 6) * 166 -> v): _*), // speed = 6
       AdvanceRateManualClock(seconds = 1),
-      CheckLastBatch(12 until 20: _*), // speed = 8
+      CheckLastBatch((12 until 20).map(v => 3000 + (v - 12) * 125 -> v): _*), // speed = 8
       AdvanceRateManualClock(seconds = 1),
       // Now we should reach full speed
-      CheckLastBatch(20 until 30: _*), // speed = 10
+      CheckLastBatch((20 until 30).map(v => 4000 + (v - 20) * 100 -> v): _*), // speed = 10
       AdvanceRateManualClock(seconds = 1),
-      CheckLastBatch(30 until 40: _*), // speed = 10
+      CheckLastBatch((30 until 40).map(v => 5000 + (v - 30) * 100 -> v): _*), // speed = 10
       AdvanceRateManualClock(seconds = 1),
-      CheckLastBatch(40 until 50: _*) // speed = 10
+      CheckLastBatch((40 until 50).map(v => 6000 + (v - 40) * 100 -> v): _*) // speed = 10
     )
   }