address comments from @zsxwing

Tyson Condie · Tyson Condie · commit 2487a7260ec1 · 2017-01-31T14:29:24.000-08:00
diff --git a/external/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetReader.scala b/external/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetReader.scala
@@ -117,7 +117,9 @@ private[kafka010] class KafkaOffsetReaderImpl(
     assert(Thread.currentThread().isInstanceOf[UninterruptibleThread])
     // Poll to get the latest assigned partitions
     consumer.poll(0)
-    consumer.assignment().asScala.toSet
+    val partitions = consumer.assignment()
+    consumer.pause(partitions)
+    partitions.asScala.toSet
   }
 
   override def fetchSpecificStartingOffsets(
diff --git a/external/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaRelation.scala b/external/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaRelation.scala
@@ -111,10 +111,10 @@ private[kafka010] class KafkaRelation(
     // Obtain TopicPartition offsets with late binding support
     kafkaOffsets match {
       case EarliestOffsets => partitions.map {
-        case tp => tp -> -2L
+        case tp => tp -> KafkaUtils.EARLIEST
       }.toMap
       case LatestOffsets => partitions.map {
-        case tp => tp -> -1L
+        case tp => tp -> KafkaUtils.LATEST
       }.toMap
       case SpecificOffsets(partitionOffsets) =>
         validateTopicPartitions(partitions, partitionOffsets)
diff --git a/external/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaSource.scala b/external/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaSource.scala
@@ -123,7 +123,7 @@ private[kafka010] class KafkaSource(
   private def fetchAndVerify(specificOffsets: Map[TopicPartition, Long]) = {
     val result = kafkaReader.fetchSpecificStartingOffsets(specificOffsets)
     specificOffsets.foreach {
-      case (tp, off) if off != -1 && off != -2 =>
+      case (tp, off) if off != KafkaUtils.LATEST && off != KafkaUtils.EARLIEST =>
         if (result(tp) != off) {
           reportDataLoss(
             s"startingOffsets for $tp was $off but consumer reset to ${result(tp)}")
diff --git a/external/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaSourceProvider.scala b/external/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaSourceProvider.scala
@@ -327,7 +327,7 @@ private[kafka010] class KafkaSourceProvider extends DataSourceRegister with Stre
         throw new IllegalArgumentException("starting relation offset can't be latest")
       case Some(json) => (SpecificOffsets(JsonUtils.partitionOffsets(json)))
         .partitionOffsets.foreach {
-          case (tp, off) if off == -1 =>
+          case (tp, off) if off == KafkaUtils.LATEST =>
             throw new IllegalArgumentException(s"startingOffsets for $tp can't be latest")
           case _ => // ignore
         }
@@ -340,7 +340,7 @@ private[kafka010] class KafkaSourceProvider extends DataSourceRegister with Stre
       case Some("latest") => // good to go
       case Some(json) => (SpecificOffsets(JsonUtils.partitionOffsets(json)))
         .partitionOffsets.foreach {
-          case (tp, off) if off == -2 =>
+          case (tp, off) if off == KafkaUtils.EARLIEST =>
             throw new IllegalArgumentException(s"ending offset for $tp can't be earliest")
           case _ => // ignore
         }
diff --git a/external/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaSourceRDD.scala b/external/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaSourceRDD.scala
@@ -140,12 +140,16 @@ private[kafka010] class KafkaSourceRDD(
         if (range.fromOffset < 0 || range.untilOffset < 0) {
           // Late bind the offset range
           val fromOffset = if (range.fromOffset < 0) {
+            assert(range.fromOffset == KafkaUtils.EARLIEST,
+              s"earliest offset does not equal ${KafkaUtils.EARLIEST}")
             consumer.rawConsumer.seekToBeginning(ju.Arrays.asList(range.topicPartition))
             consumer.rawConsumer.position(range.topicPartition)
           } else {
             range.fromOffset
           }
           val untilOffset = if (range.untilOffset < 0) {
+            assert(range.fromOffset == KafkaUtils.LATEST,
+              s"latest offset does not equal ${KafkaUtils.LATEST}")
             consumer.rawConsumer.seekToEnd(ju.Arrays.asList(range.topicPartition))
             consumer.rawConsumer.position(range.topicPartition)
           } else {
diff --git a/external/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaUtils.scala b/external/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaUtils.scala
@@ -22,6 +22,10 @@ import org.apache.spark.scheduler.ExecutorCacheTaskLocation
 
 private[kafka010] object KafkaUtils {
 
+  // Used to denote unbounded offset positions
+  val LATEST = -1L
+  val EARLIEST = -2L
+
   def getSortedExecutorList(sc: SparkContext): Array[String] = {
     val bm = sc.env.blockManager
     bm.master.getPeers(bm.blockManagerId).toArray
diff --git a/external/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaTestUtils.scala b/external/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaTestUtils.scala
@@ -50,7 +50,7 @@ import org.apache.spark.SparkConf
  *
  * The reason to put Kafka test utility class in src is to test Python related Kafka APIs.
  */
-class KafkaTestUtils(withBrokerProps: Option[Map[String, Object]] = None) extends Logging {
+class KafkaTestUtils(withBrokerProps: Map[String, Object] = Map.empty) extends Logging {
 
   // Zookeeper related configurations
   private val zkHost = "localhost"
@@ -255,15 +255,15 @@ class KafkaTestUtils(withBrokerProps: Option[Map[String, Object]] = None) extend
 
   def getEarliestOffsets(topics: Set[String]): Map[TopicPartition, Long] = {
     val kc = new KafkaConsumer[String, String](consumerConfiguration)
-    logInfo("Created consumer to get latest offsets")
+    logInfo("Created consumer to get earliest offsets")
     kc.subscribe(topics.asJavaCollection)
     kc.poll(0)
     val partitions = kc.assignment()
     kc.pause(partitions)
     kc.seekToBeginning(partitions)
     val offsets = partitions.asScala.map(p => p -> kc.position(p)).toMap
     kc.close()
-    logInfo("Closed consumer to get latest offsets")
+    logInfo("Closed consumer to get earliest offsets")
     offsets
   }
 
@@ -292,11 +292,7 @@ class KafkaTestUtils(withBrokerProps: Option[Map[String, Object]] = None) extend
     props.put("log.flush.interval.messages", "1")
     props.put("replica.socket.timeout.ms", "1500")
     props.put("delete.topic.enable", "true")
-    withBrokerProps.map { p =>
-      p.foreach {
-        case (key, value) => props.put(key, value)
-      }
-    }
+    props.putAll(withBrokerProps.asJava)
     props
   }
 
