Rollback KafkaMicroBatchStream reanmes

Author: gaborgsomogyi

external/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaMicroBatchStream.scala

@@ -56,7 +56,7 @@ import org.apache.spark.util.UninterruptibleThread
  * and not use wrong broker addresses.
  */
 private[kafka010] class KafkaMicroBatchStream(
-    private val offsetReader: KafkaOffsetReader,
+    private val kafkaOffsetReader: KafkaOffsetReader,
     executorKafkaParams: ju.Map[String, Object],
     options: CaseInsensitiveStringMap,
     metadataPath: String,
@@ -85,7 +85,7 @@ private[kafka010] class KafkaMicroBatchStream(
 
   override def latestOffset(start: Offset): Offset = {
     val startPartitionOffsets = start.asInstanceOf[KafkaSourceOffset].partitionToOffsets
-    val latestPartitionOffsets = offsetReader.fetchLatestOffsets(Some(startPartitionOffsets))
+    val latestPartitionOffsets = kafkaOffsetReader.fetchLatestOffsets(Some(startPartitionOffsets))
     endPartitionOffsets = KafkaSourceOffset(maxOffsetsPerTrigger.map { maxOffsets =>
       rateLimit(maxOffsets, startPartitionOffsets, latestPartitionOffsets)
     }.getOrElse {
@@ -100,7 +100,7 @@ private[kafka010] class KafkaMicroBatchStream(
 
     // Find the new partitions, and get their earliest offsets
     val newPartitions = endPartitionOffsets.keySet.diff(startPartitionOffsets.keySet)
-    val newPartitionInitialOffsets = offsetReader.fetchEarliestOffsets(newPartitions.toSeq)
+    val newPartitionInitialOffsets = kafkaOffsetReader.fetchEarliestOffsets(newPartitions.toSeq)
     if (newPartitionInitialOffsets.keySet != newPartitions) {
       // We cannot get from offsets for some partitions. It means they got deleted.
       val deletedPartitions = newPartitions.diff(newPartitionInitialOffsets.keySet)
@@ -117,7 +117,7 @@ private[kafka010] class KafkaMicroBatchStream(
     val deletedPartitions = startPartitionOffsets.keySet.diff(endPartitionOffsets.keySet)
     if (deletedPartitions.nonEmpty) {
       val message =
-        if (offsetReader.driverKafkaParams.containsKey(ConsumerConfig.GROUP_ID_CONFIG)) {
+        if (kafkaOffsetReader.driverKafkaParams.containsKey(ConsumerConfig.GROUP_ID_CONFIG)) {
           s"$deletedPartitions are gone. ${KafkaSourceProvider.CUSTOM_GROUP_ID_ERROR_MESSAGE}"
         } else {
           s"$deletedPartitions are gone. Some data may have been missed."
@@ -172,10 +172,10 @@ private[kafka010] class KafkaMicroBatchStream(
   override def commit(end: Offset): Unit = {}
 
   override def stop(): Unit = {
-    offsetReader.close()
+    kafkaOffsetReader.close()
   }
 
-  override def toString(): String = s"KafkaV2[$offsetReader]"
+  override def toString(): String = s"KafkaV2[$kafkaOffsetReader]"
 
   /**
    * Read initial partition offsets from the checkpoint, or decide the offsets and write them to
@@ -195,11 +195,11 @@ private[kafka010] class KafkaMicroBatchStream(
     metadataLog.get(0).getOrElse {
       val offsets = startingOffsets match {
         case EarliestOffsetRangeLimit =>
-          KafkaSourceOffset(offsetReader.fetchEarliestOffsets())
+          KafkaSourceOffset(kafkaOffsetReader.fetchEarliestOffsets())
         case LatestOffsetRangeLimit =>
-          KafkaSourceOffset(offsetReader.fetchLatestOffsets(None))
+          KafkaSourceOffset(kafkaOffsetReader.fetchLatestOffsets(None))
         case SpecificOffsetRangeLimit(p) =>
-          offsetReader.fetchSpecificOffsets(p, reportDataLoss)
+          kafkaOffsetReader.fetchSpecificOffsets(p, reportDataLoss)
       }
       metadataLog.add(0, offsets)
       logInfo(s"Initial offsets: $offsets")
@@ -212,7 +212,7 @@ private[kafka010] class KafkaMicroBatchStream(
       limit: Long,
       from: PartitionOffsetMap,
       until: PartitionOffsetMap): PartitionOffsetMap = {
-    val fromNew = offsetReader.fetchEarliestOffsets(until.keySet.diff(from.keySet).toSeq)
+    val fromNew = kafkaOffsetReader.fetchEarliestOffsets(until.keySet.diff(from.keySet).toSeq)
     val sizes = until.flatMap {
       case (tp, end) =>
         // If begin isn't defined, something's wrong, but let alert logic in getBatch handle it

external/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaSourceProviderSuite.scala

@@ -33,7 +33,6 @@ class KafkaSourceProviderSuite extends SparkFunSuite with PrivateMethodTester {
   private val expected = "1111"
   private val pollTimeoutMsMethod = PrivateMethod[Long]('pollTimeoutMs)
   private val maxOffsetsPerTriggerMethod = PrivateMethod[Option[Long]]('maxOffsetsPerTrigger)
-  private val offsetReaderMethod = PrivateMethod[KafkaOffsetReader]('offsetReader)
   private val maxOffsetFetchAttemptsMethod = PrivateMethod[Int]('maxOffsetFetchAttempts)
   private val offsetFetchAttemptIntervalMsMethod =
     PrivateMethod[Long]('offsetFetchAttemptIntervalMs)
@@ -50,6 +49,8 @@ class KafkaSourceProviderSuite extends SparkFunSuite with PrivateMethodTester {
   }
 
   test("micro-batch mode - options should be handled as case-insensitive") {
+    val offsetReaderMethod = PrivateMethod[KafkaOffsetReader]('kafkaOffsetReader)
+
     verifyFieldsInMicroBatchStream(KafkaSourceProvider.CONSUMER_POLL_TIMEOUT, expected, stream => {
       assert(expected.toLong === getField(stream, pollTimeoutMsMethod))
     })
@@ -68,6 +69,8 @@ class KafkaSourceProviderSuite extends SparkFunSuite with PrivateMethodTester {
   }
 
   test("continuous mode - options should be handled as case-insensitive") {
+    val offsetReaderMethod = PrivateMethod[KafkaOffsetReader]('offsetReader)
+
     verifyFieldsInContinuousStream(KafkaSourceProvider.CONSUMER_POLL_TIMEOUT, expected, stream => {
       assert(expected.toLong === getField(stream, pollTimeoutMsMethod))
     })