Bug 32014 - Spark Consumer fails with java.lang.AssertionError (apach…

…e#326) (apache#336)

* MapR [32014] Spark Consumer fails with java.lang.AssertionError

external/kafka-0-9/src/main/scala/org/apache/spark/streaming/kafka09/CachedKafkaConsumer.scala

@@ -51,8 +51,6 @@ class CachedKafkaConsumer[K, V] private(
     c
   }
 
-  val isStreams = topic.startsWith("/") && topic.contains(":")
-
   // TODO if the buffer was kept around as a random-access structure,
   // could possibly optimize re-calculating of an RDD in the same batch
   protected var buffer = ju.Collections.emptyList[ConsumerRecord[K, V]]().iterator
@@ -80,8 +78,9 @@ class CachedKafkaConsumer[K, V] private(
 
     nextOffset = offset + 1
 
-    if (record.offset() == KafkaUtils.eofOffset && isStreams && buffer.hasNext) {
-      buffer.next()
+    if (record.offset() == KafkaUtils.eofOffset) {
+      log.debug("EOF message is received")
+      if (buffer.hasNext) buffer.next() else null
     } else {
       record
     }

external/kafka-0-9/src/main/scala/org/apache/spark/streaming/kafka09/KafkaRDD.scala

@@ -209,20 +209,47 @@ private[spark] class KafkaRDD[K, V](
     }
 
     var requestOffset = part.fromOffset
+    var currentRecord: ConsumerRecord[K, V] = _
 
     def closeIfNeeded(): Unit = {
       if (!useConsumerCache && consumer != null) {
         consumer.close
       }
     }
 
-    override def hasNext(): Boolean = requestOffset < part.untilOffset
+    private def setNext() = {
+      if (currentRecord == null) {
+        currentRecord = consumer.get(requestOffset, pollTimeout)
+
+        requestOffset =
+          if (currentRecord == null) {
+            logInfo(s"Skipping offsets from $requestOffset to ${part.untilOffset}")
+            part.untilOffset
+          } else {
+            currentRecord.offset() + 1
+          }
+      }
+    }
+
+    override def hasNext(): Boolean = {
+      if (currentRecord != null) {
+        true
+      } else if (requestOffset < part.untilOffset) {
+        setNext()
+
+        currentRecord != null
+      } else {
+        false
+      }
+    }
 
     override def next(): ConsumerRecord[K, V] = {
       assert(hasNext(), "Can't call getNext() once untilOffset has been reached")
-      val r = consumer.get(requestOffset, pollTimeout)
-      requestOffset += 1
-      r
+
+      val rec = currentRecord
+      currentRecord = null
+
+      rec
     }
   }
 }