[MAPR-26060] Fixed case when mapr-streams make gaps in offsets (apache#97)

Author: rsotn-mapr

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

@@ -19,10 +19,11 @@ package org.apache.spark.streaming.kafka09
 
 import java.{ util => ju }
 
-import org.apache.kafka.clients.consumer.{ ConsumerConfig, ConsumerRecord, KafkaConsumer }
-import org.apache.kafka.common.{ KafkaException, TopicPartition }
+import collection.JavaConverters._
+
+import org.apache.kafka.clients.consumer.{ConsumerConfig, ConsumerRecord, KafkaConsumer}
+import org.apache.kafka.common.{KafkaException, TopicPartition}
 
-import org.apache.spark.SparkConf
 import org.apache.spark.internal.Logging
 
 
@@ -73,9 +74,12 @@ class CachedKafkaConsumer[K, V] private(
     if (!buffer.hasNext()) { poll(timeout) }
     assert(buffer.hasNext(),
       s"Failed to get records for $groupId $topic $partition $offset after polling for $timeout")
-    var record = buffer.next()
+    val record = buffer.next()
 
-    if (record.offset != offset) {
+    nextOffset = offset + 1
+    record
+// Offsets in MapR-streams can contains gaps
+/*    if (record.offset < offset) {
       logInfo(s"Buffer miss for $groupId $topic $partition $offset")
       seek(offset)
       poll(timeout)
@@ -88,6 +92,7 @@ class CachedKafkaConsumer[K, V] private(
 
     nextOffset = offset + 1
     record
+        */
   }
 
   private def seek(offset: Long): Unit = {
@@ -98,8 +103,15 @@ class CachedKafkaConsumer[K, V] private(
   private def poll(timeout: Long): Unit = {
     val p = consumer.poll(timeout)
     val r = p.records(topicPartition)
-    logDebug(s"Polled ${p.partitions()}  ${r.size}")
-    buffer = r.iterator
+
+    val preparedRecords = if (r.size > 1) {
+      r.asScala.filter(_.offset() > 0).asJava
+    } else {
+      r
+    }
+
+    logDebug(s"Polled ${p.partitions()}  ${preparedRecords.size}")
+    buffer = preparedRecords.iterator()
   }
 }
 

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

@@ -30,6 +30,9 @@ import org.apache.spark.partial.{BoundedDouble, PartialResult}
 import org.apache.spark.rdd.RDD
 import org.apache.spark.scheduler.ExecutorCacheTaskLocation
 import org.apache.spark.storage.StorageLevel
+import org.apache.spark.util.NextIterator
+
+import scala.annotation.tailrec
 
 /**
  * A batch-oriented interface for consuming from Kafka.
@@ -83,11 +86,11 @@ private[spark] class KafkaRDD[K, V](
 
   override def getPartitions: Array[Partition] = {
     offsetRanges.zipWithIndex.map { case (o, i) =>
-        new KafkaRDDPartition(i, o.topic, o.partition, o.fromOffset, o.untilOffset)
+      new KafkaRDDPartition(i, o.topic, o.partition, o.fromOffset, o.untilOffset)
     }.toArray
   }
 
-  override def count(): Long = offsetRanges.map(_.count).sum
+  //  override def count(): Long = offsetRanges.map(_.count).sum
 
   override def countApprox(
       timeout: Long,
@@ -193,7 +196,7 @@ private[spark] class KafkaRDD[K, V](
    */
   private class KafkaRDDIterator(
       part: KafkaRDDPartition,
-      context: TaskContext) extends Iterator[ConsumerRecord[K, V]] {
+      context: TaskContext) extends NextIterator[ConsumerRecord[K, V]] {
 
     logInfo(s"Computing topic ${part.topic}, partition ${part.partition} " +
       s"offsets ${part.fromOffset} -> ${part.untilOffset}")
@@ -215,19 +218,41 @@ private[spark] class KafkaRDD[K, V](
 
     var requestOffset = part.fromOffset
 
-    def closeIfNeeded(): Unit = {
+    override def close(): Unit = {
       if (!useConsumerCache && consumer != null) {
         consumer.close
       }
     }
 
-    override def hasNext(): Boolean = requestOffset < part.untilOffset
+    //    override def hasNext(): Boolean = requestOffset < part.untilOffset
+    override def getNext(): ConsumerRecord[K, V] = {
+
+      @tailrec
+      def skipGapsAndGetNext: ConsumerRecord[K, V] = {
+        if (requestOffset < part.untilOffset) {
+          val r = consumer.get(requestOffset, pollTimeout)
+
+          requestOffset = if (r.offset() == 0) {part.untilOffset} else {r.offset() + 1}
+
+          if (null == r && r.offset() == 0) {
+            skipGapsAndGetNext
+          } else {
+            r
+          }
+        } else {
+          finished = true
+          null.asInstanceOf[ConsumerRecord[K, V]]
+        }
+      }
 
-    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
+      skipGapsAndGetNext
     }
+
+    //    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
+    //    }
   }
 }

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

@@ -206,6 +206,9 @@ object KafkaUtils extends Logging {
   * Tweak kafka params to prevent issues on executors
   */
   private[kafka09] def fixKafkaParams(kafkaParams: ju.HashMap[String, Object]): Unit = {
+    logWarning(s"overriding ${ConsumerConfig.STREAMS_ZEROOFFSET_RECORD_ON_EOF_CONFIG} to true")
+    kafkaParams.put(ConsumerConfig.STREAMS_ZEROOFFSET_RECORD_ON_EOF_CONFIG, true: java.lang.Boolean)
+
     logWarning(s"overriding ${ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG} to false for executor")
     kafkaParams.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, false: java.lang.Boolean)