The speed should be always <= tuplesPerSecond

Author: zsxwing

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

@@ -110,6 +110,7 @@ class RateStreamSource(
     useManualClock: Boolean) extends Source with Logging {
 
   import RateSourceProvider._
+  import RateStreamSource._
 
   val clock = if (useManualClock) new ManualClock else new SystemClock
 
@@ -183,15 +184,8 @@ class RateStreamSource(
     if (lastTimeMs < TimeUnit.SECONDS.toMillis(endSeconds) + startTimeMs) {
       lastTimeMs = TimeUnit.SECONDS.toMillis(endSeconds) + startTimeMs
     }
-    val (rangeStart, rangeEnd) = if (rampUpTimeSeconds > endSeconds) {
-      (math.rint(tuplesPerSecond * (startSeconds * 1.0 / rampUpTimeSeconds)).toLong * startSeconds,
-        math.rint(tuplesPerSecond * (endSeconds * 1.0 / rampUpTimeSeconds)).toLong * endSeconds)
-    } else if (startSeconds < rampUpTimeSeconds) {
-      (math.rint(tuplesPerSecond * (startSeconds * 1.0 / rampUpTimeSeconds)).toLong * startSeconds,
-        endSeconds * tuplesPerSecond)
-    } else {
-      (startSeconds * tuplesPerSecond, endSeconds * tuplesPerSecond)
-    }
+    val rangeStart = valueAtSecond(startSeconds, tuplesPerSecond, rampUpTimeSeconds)
+    val rangeEnd = valueAtSecond(endSeconds, tuplesPerSecond, rampUpTimeSeconds)
     logDebug(s"startSeconds: $startSeconds, endSeconds: $endSeconds, " +
       s"rangeStart: $rangeStart, rangeEnd: $rangeEnd")
     val localStartTimeMs = startTimeMs
@@ -206,3 +200,30 @@ class RateStreamSource(
 
   override def stop(): Unit = {}
 }
+
+object RateStreamSource {
+
+  /** Calculate the end value we will emit at the time `seconds`. */
+  def valueAtSecond(seconds: Long, tuplesPerSecond: Long, rampUpTimeSeconds: Long): Long = {
+    // E.g., rampUpTimeSeconds = 4, tuplesPerSecond = 10
+    // Then speedDeltaPerSecond = 2
+    //
+    // seconds   = 0 1 2  3  4  5  6
+    // speed     = 0 2 4  6  8 10 10 (speedDeltaPerSecond * seconds)
+    // end value = 0 2 6 12 20 30 40 (0 + speedDeltaPerSecond * seconds) * (seconds + 1) / 2
+    val speedDeltaPerSecond = tuplesPerSecond / (rampUpTimeSeconds + 1)
+    if (seconds <= rampUpTimeSeconds) {
+      // Calculate "(0 + speedDeltaPerSecond * seconds) * (seconds + 1) / 2" in a special way to
+      // avoid overflow
+      if (seconds % 2 == 1) {
+        (seconds + 1) / 2 * speedDeltaPerSecond * seconds
+      } else {
+        seconds / 2 * speedDeltaPerSecond * (seconds + 1)
+      }
+    } else {
+      // rampUpPart is just a special case of the above formula: rampUpTimeSeconds == seconds
+      val rampUpPart = valueAtSecond(rampUpTimeSeconds, tuplesPerSecond, rampUpTimeSeconds)
+      rampUpPart + (seconds - rampUpTimeSeconds) * tuplesPerSecond
+    }
+  }
+}

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

@@ -64,27 +64,46 @@ class RateSourceSuite extends StreamTest {
     )
   }
 
+  test("valueAtSecond") {
+    import RateStreamSource._
+
+    assert(valueAtSecond(seconds = 0, tuplesPerSecond = 5, rampUpTimeSeconds = 2) === 0)
+    assert(valueAtSecond(seconds = 1, tuplesPerSecond = 5, rampUpTimeSeconds = 2) === 1)
+    assert(valueAtSecond(seconds = 2, tuplesPerSecond = 5, rampUpTimeSeconds = 2) === 3)
+    assert(valueAtSecond(seconds = 3, tuplesPerSecond = 5, rampUpTimeSeconds = 2) === 8)
+
+    assert(valueAtSecond(seconds = 0, tuplesPerSecond = 10, rampUpTimeSeconds = 4) === 0)
+    assert(valueAtSecond(seconds = 1, tuplesPerSecond = 10, rampUpTimeSeconds = 4) === 2)
+    assert(valueAtSecond(seconds = 2, tuplesPerSecond = 10, rampUpTimeSeconds = 4) === 6)
+    assert(valueAtSecond(seconds = 3, tuplesPerSecond = 10, rampUpTimeSeconds = 4) === 12)
+    assert(valueAtSecond(seconds = 4, tuplesPerSecond = 10, rampUpTimeSeconds = 4) === 20)
+    assert(valueAtSecond(seconds = 5, tuplesPerSecond = 10, rampUpTimeSeconds = 4) === 30)
+  }
+
   test("rampUpTimeSeconds") {
     val input = spark.readStream
       .format("rate")
       .option("tuplesPerSecond", "10")
-      .option("rampUpTimeSeconds", "5")
+      .option("rampUpTimeSeconds", "4")
       .option("useManualClock", "true")
       .load()
       .select($"value")
     testStream(input)(
       AdvanceRateManualClock(seconds = 1),
-      CheckLastBatch((0 until 2): _*),
+      CheckLastBatch(0 until 2: _*), // speed = 2
+      AdvanceRateManualClock(seconds = 1),
+      CheckLastBatch(2 until 6: _*), // speed = 4
       AdvanceRateManualClock(seconds = 1),
-      CheckLastBatch((2 until 8): _*),
+      CheckLastBatch(6 until 12: _*), // speed = 6
       AdvanceRateManualClock(seconds = 1),
-      CheckLastBatch((8 until 18): _*),
+      CheckLastBatch(12 until 20: _*), // speed = 8
       AdvanceRateManualClock(seconds = 1),
-      CheckLastBatch((18 until 32): _*),
+      // Now we should reach full speed
+      CheckLastBatch(20 until 30: _*), // speed = 10
       AdvanceRateManualClock(seconds = 1),
-      CheckLastBatch((32 until 50): _*),
+      CheckLastBatch(30 until 40: _*), // speed = 10
       AdvanceRateManualClock(seconds = 1),
-      CheckLastBatch((50 until 60): _*)
+      CheckLastBatch(40 until 50: _*) // speed = 10
     )
   }