additional testing, fix fencepost error

apache · Dec 25, 2014 · cac63ee · cac63ee
1 parent 37d3053
commit cac63ee
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 6 deletions.
diff --git a/external/kafka/src/main/scala/org/apache/spark/rdd/kafka/KafkaRDD.scala b/external/kafka/src/main/scala/org/apache/spark/rdd/kafka/KafkaRDD.scala
@@ -132,12 +132,14 @@ class KafkaRDD[
             null.asInstanceOf[R]
           } else {
             val item = iter.next
-            if (item.offset > part.untilOffset) {
+            if (item.offset >= part.untilOffset) {
               finished = true
+              null.asInstanceOf[R]
+            } else {
+              requestOffset = item.nextOffset
+              messageHandler(new MessageAndMetadata(
+                part.topic, part.partition, item.message, item.offset, keyDecoder, valueDecoder))
             }
-            requestOffset = item.nextOffset
-            messageHandler(new MessageAndMetadata(
-              part.topic, part.partition, item.message, item.offset, keyDecoder, valueDecoder))
           }
         }
       }

diff --git a/external/kafka/src/test/scala/org/apache/spark/rdd/kafka/KafkaRDDSuite.scala b/external/kafka/src/test/scala/org/apache/spark/rdd/kafka/KafkaRDDSuite.scala
@@ -55,13 +55,21 @@ class KafkaRDDSuite extends KafkaStreamSuiteBase with BeforeAndAfter {
 
     val rdd = getRdd(kc, Set(topic))
     assert(rdd.isDefined)
-    assert(rdd.get.countByValue.size === sent.size)
+    assert(rdd.get.count === sent.values.sum)
 
     kc.setConsumerOffsets(kafkaParams("group.id"), rdd.get.untilOffsets)
 
     val rdd2 = getRdd(kc, Set(topic))
+    val sent2 = Map("d" -> 1)
+    produceAndSendMessage(topic, sent2)
     assert(rdd2.isDefined)
     assert(rdd2.get.count === 0)
+
+    val rdd3 = getRdd(kc, Set(topic))
+    produceAndSendMessage(topic, Map("extra" -> 22))
+    assert(rdd3.isDefined)
+    assert(rdd3.get.count === sent2.values.sum)
+
   }
 
   private def getRdd(kc: KafkaCluster, topics: Set[String]) = {
@@ -73,7 +81,7 @@ class KafkaRDDSuite extends KafkaStreamSuiteBase with BeforeAndAfter {
       until <- kc.getLatestLeaderOffsets(topicPartitions).right.toOption
     } yield {
       new KafkaRDD[String, String, StringDecoder, StringDecoder, String](
-        sc, kc.kafkaParams, from, until, mmd => mmd.message)
+        sc, kc.kafkaParams, from, until, mmd => s"${mmd.offset} ${mmd.message}")
     }
   }
 }