Change to while loop

marmbrus · Nov 12, 2014 · e5e21c1 · e5e21c1
1 parent ea873e4
commit e5e21c1
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/external/kafka/src/main/scala/org/apache/spark/streaming/kafka/KafkaInputDStream.scala b/external/kafka/src/main/scala/org/apache/spark/streaming/kafka/KafkaInputDStream.scala
@@ -128,7 +128,9 @@ class KafkaReceiver[
     def run() {
       logInfo("Starting MessageHandler.")
       try {
-        for (msgAndMetadata <- stream) {
+        val streamIterator = stream.iterator()
+        while (streamIterator.hasNext()) {
+          val msgAndMetadata = streamIterator.next()
           store((msgAndMetadata.key, msgAndMetadata.message))
         }
       } catch {

diff --git a/external/kafka/src/main/scala/org/apache/spark/streaming/kafka/ReliableKafkaReceiver.scala b/external/kafka/src/main/scala/org/apache/spark/streaming/kafka/ReliableKafkaReceiver.scala
@@ -206,7 +206,9 @@ class ReliableKafkaReceiver[
     override def run(): Unit = {
       logInfo(s"Starting message process thread ${Thread.currentThread().getId}.")
       try {
-        for (msgAndMetadata <- stream) {
+        val streamIterator = stream.iterator()
+        while (streamIterator.hasNext()) {
+          val msgAndMetadata = streamIterator.next()
           val topicAndPartition = TopicAndPartition(
             msgAndMetadata.topic, msgAndMetadata.partition)
           blockGenerator.synchronized {