[SPARK-17813][SQL][KAFKA] Maximum data per trigger

[koeninger]

What changes were proposed in this pull request?

maxOffsetsPerTrigger option for rate limiting, proportionally based on volume of different topicpartitions.

How was this patch tested?

Added unit test

[SparkQA]

Test build #67113 has finished for PR 15527 at commit 6c8d459.

• This patch passes all tests.
• This patch merges cleanly.
• This patch adds no public classes.

[SparkQA]

Test build #67455 has finished for PR 15527 at commit 76ecafc.

• This patch fails Spark unit tests.
• This patch merges cleanly.
• This patch adds no public classes.

[SparkQA]

Test build #67459 has finished for PR 15527 at commit fde4e33.

• This patch passes all tests.
• This patch merges cleanly.
• This patch adds no public classes.

[zsxwing]

Thanks for your work. The major issue is currentPartitionOffsets should be set when recovering from failure.

[zsxwing]

currentPartitionOffsets should be set to untilPartitionOffsets if it's empty. It means recovery from a failure.

[koeninger]

Shouldn't it be set to the highest available offset in the streaming metadata log, not the highest available offset in kafka?

[zsxwing]

currentPartitionOffsets is the last processed offsets. Right? When recovering from the failure, getBatch will be called firstly, then getOffset.

[koeninger]

Ok, I see what you were saying, thought you were talking about resetting it in getOffset, not the end passed to getBatch.

[zsxwing]

There is a race condition here. The batch may be still running. I figured out the following codes to cover the recovery and fix the race condition finally.

  test("maxOffsetsPerTrigger") {
    val topic = newTopic()
    testUtils.createTopic(topic, partitions = 3)
    testUtils.sendMessages(topic, (100 to 200).map(_.toString).toArray, Some(0))
    testUtils.sendMessages(topic, (10 to 20).map(_.toString).toArray, Some(1))
    testUtils.sendMessages(topic, Array("1"), Some(2))

    val reader = spark
      .readStream
      .format("kafka")
      .option("kafka.bootstrap.servers", testUtils.brokerAddress)
      .option("kafka.metadata.max.age.ms", "1")
      .option("maxOffsetsPerTrigger", 10)
      .option("subscribe", topic)
      .option("startingOffsets", "earliest")
    val kafka = reader.load()
      .selectExpr("CAST(key AS STRING)", "CAST(value AS STRING)")
      .as[(String, String)]
    val mapped: org.apache.spark.sql.Dataset[_] = kafka.map(kv => kv._2.toInt)

    val clock = new StreamManualClock

    val waitUntilBatchProcessed = AssertOnQuery { q =>
      eventually(Timeout(streamingTimeout)) {
        if (!q.exception.isDefined) {
          assert(clock.isStreamWaitingAt(clock.getTimeMillis()))
        }
      }
      if (q.exception.isDefined) {
        throw q.exception.get
      }
      true
    }

    testStream(mapped)(
      StartStream(ProcessingTime(100), clock),
      waitUntilBatchProcessed,
      // 1 from smallest, 1 from middle, 8 from biggest
      CheckAnswer(1, 10, 100, 101, 102, 103, 104, 105, 106, 107),
      AdvanceManualClock(100),
      waitUntilBatchProcessed,
      // smallest now empty, 1 more from middle, 9 more from biggest
      CheckAnswer(1, 10, 100, 101, 102, 103, 104, 105, 106, 107,
        11, 108, 109, 110, 111, 112, 113, 114, 115, 116
      ),
      StopStream,
      StartStream(ProcessingTime(100), clock),
      waitUntilBatchProcessed,
      AdvanceManualClock(100),
      waitUntilBatchProcessed,
      // smallest now empty, 1 more from middle, 9 more from biggest
      CheckAnswer(1, 10, 100, 101, 102, 103, 104, 105, 106, 107,
        11, 108, 109, 110, 111, 112, 113, 114, 115, 116,
        12, 117, 118, 119, 120, 121, 122, 123, 124, 125
      ),
      AdvanceManualClock(100),
      waitUntilBatchProcessed,
      // smallest now empty, 1 more from middle, 9 more from biggest
      CheckAnswer(1, 10, 100, 101, 102, 103, 104, 105, 106, 107,
        11, 108, 109, 110, 111, 112, 113, 114, 115, 116,
        12, 117, 118, 119, 120, 121, 122, 123, 124, 125,
        13, 126, 127, 128, 129, 130, 131, 132, 133, 134
      )
    )
  }

This test fails now because of an issue being fixed in #14553.

[zsxwing]

FYI, #14553 got merged.

[SparkQA]

Test build #67616 has finished for PR 15527 at commit 6a7ff24.

• This patch passes all tests.
• This patch merges cleanly.
• This patch adds no public classes.

[zsxwing]

LGTM except some style nits.

[zsxwing]

nit: use 2 spaces

[zsxwing]

nit: use 2 spaces

[SparkQA]

Test build #67645 has finished for PR 15527 at commit 5e4b468.

• This patch passes all tests.
• This patch merges cleanly.
• This patch adds no public classes.

[zsxwing]

LGTM. Thanks! Merging to master and 2.0.