[SPARK-25233][STREAMING] Give the user the option of specifying a minimum message per partition per batch when using kafka direct API with backpressure

After SPARK-18371, it is guaranteed that there would be at least one message per partition per batch using direct kafka API when new messages exist in the topics. This change will give the user the option of setting the minimum instead of just a hard coded 1 limit
The related unit test is updated and some internal tests verified that the topic partitions with new messages will be progressed by the specified minimum.

Author: Reza Safi <[email protected]>

Closes apache#22223 from rezasafi/streaminglag.

Author: Reza Safi

docs/configuration.md

@@ -1997,6 +1997,14 @@ showDF(properties, numRows = 200, truncate = FALSE)
     for more details.
   </td>
 </tr>
+<tr>
+    <td><code>spark.streaming.kafka.minRatePerPartition</code></td>
+    <td>1</td>
+    <td>
+      Minimum rate (number of records per second) at which data will be read from each Kafka
+      partition when using the new Kafka direct stream API.
+    </td>
+</tr>
 <tr>
   <td><code>spark.streaming.kafka.maxRetries</code></td>
   <td>1</td>

external/kafka-0-10/src/main/scala/org/apache/spark/streaming/kafka010/DirectKafkaInputDStream.scala

@@ -154,7 +154,8 @@ private[spark] class DirectKafkaInputDStream[K, V](
     if (effectiveRateLimitPerPartition.values.sum > 0) {
       val secsPerBatch = context.graph.batchDuration.milliseconds.toDouble / 1000
       Some(effectiveRateLimitPerPartition.map {
-        case (tp, limit) => tp -> Math.max((secsPerBatch * limit).toLong, 1L)
+        case (tp, limit) => tp -> Math.max((secsPerBatch * limit).toLong,
+          ppc.minRatePerPartition(tp))
       })
     } else {
       None

external/kafka-0-10/src/main/scala/org/apache/spark/streaming/kafka010/PerPartitionConfig.scala

@@ -34,6 +34,7 @@ abstract class PerPartitionConfig extends Serializable {
    *  from each Kafka partition.
    */
   def maxRatePerPartition(topicPartition: TopicPartition): Long
+  def minRatePerPartition(topicPartition: TopicPartition): Long = 1
 }
 
 /**
@@ -42,6 +43,8 @@ abstract class PerPartitionConfig extends Serializable {
 private class DefaultPerPartitionConfig(conf: SparkConf)
     extends PerPartitionConfig {
   val maxRate = conf.getLong("spark.streaming.kafka.maxRatePerPartition", 0)
+  val minRate = conf.getLong("spark.streaming.kafka.minRatePerPartition", 1)
 
   def maxRatePerPartition(topicPartition: TopicPartition): Long = maxRate
+  override def minRatePerPartition(topicPartition: TopicPartition): Long = minRate
 }

external/kafka-0-10/src/test/scala/org/apache/spark/streaming/kafka010/DirectKafkaStreamSuite.scala

@@ -664,7 +664,8 @@ class DirectKafkaStreamSuite
     kafkaStream.stop()
   }
 
-  test("maxMessagesPerPartition with zero offset and rate equal to one") {
+  test("maxMessagesPerPartition with zero offset and rate equal to the specified" +
+    " minimum with default 1") {
     val topic = "backpressure"
     val kafkaParams = getKafkaParams()
     val batchIntervalMilliseconds = 60000
@@ -674,6 +675,8 @@ class DirectKafkaStreamSuite
       .setMaster("local[1]")
       .setAppName(this.getClass.getSimpleName)
       .set("spark.streaming.kafka.maxRatePerPartition", "100")
+      .set("spark.streaming.kafka.minRatePerPartition", "5")
+
 
     // Setup the streaming context
     ssc = new StreamingContext(sparkConf, Milliseconds(batchIntervalMilliseconds))
@@ -704,12 +707,13 @@ class DirectKafkaStreamSuite
     )
     val result = kafkaStream.maxMessagesPerPartition(offsets)
     val expected = Map(
-      new TopicPartition(topic, 0) -> 1L,
+      new TopicPartition(topic, 0) -> 5L,
       new TopicPartition(topic, 1) -> 10L,
       new TopicPartition(topic, 2) -> 20L,
       new TopicPartition(topic, 3) -> 30L
     )
-    assert(result.contains(expected), s"Number of messages per partition must be at least 1")
+    assert(result.contains(expected), s"Number of messages per partition must be at least equal" +
+      s" to the specified minimum")
   }
 
   /** Get the generated offset ranges from the DirectKafkaStream */