Batch oriented kafka rdd, WIP. todo: cluster metadata / finding leader

Author: koeninger

external/kafka/pom.xml

@@ -44,7 +44,7 @@
     <dependency>
       <groupId>org.apache.kafka</groupId>
       <artifactId>kafka_${scala.binary.version}</artifactId>
-      <version>0.8.0</version>
+      <version>0.8.1.1</version>
       <exclusions>
         <exclusion>
           <groupId>com.sun.jmx</groupId>

external/kafka/src/main/scala/org/apache/spark/rdd/kafka/KafkaRDD.scala

@@ -0,0 +1,126 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.rdd.kafka
+
+import scala.reflect.{classTag, ClassTag}
+
+import org.apache.spark.{Logging, Partition, SparkContext, TaskContext}
+import org.apache.spark.rdd.RDD
+import org.apache.spark.util.NextIterator
+
+import java.util.Properties
+import kafka.api.FetchRequestBuilder
+import kafka.common.{ErrorMapping, TopicAndPartition}
+import kafka.consumer.{ConsumerConfig, SimpleConsumer}
+import kafka.message.{MessageAndMetadata, MessageAndOffset}
+import kafka.serializer.Decoder
+import kafka.utils.VerifiableProperties
+
+private[spark] case class KafkaRDDPartition(
+  override val index: Int,
+  topic: String,
+  partition: Int,
+  fromOffset: Long,
+  untilOffset: Long
+) extends Partition
+
+/** A batch-oriented interface to Kafka.
+  * Each given Kafka topic/partition corresponds to an RDD partition.
+  * Starting and ending offsets are specified in advance, so that you can control exactly-once semantics.
+  * For an easy interface to Kafka-managed offsets, see {@link org.apache.spark.rdd.kafka.KafkaCluster}
+  * @param kafkaParams Kafka <a href="http://kafka.apache.org/documentation.html#configuration">configuration parameters</a>.
+  *   Requires "metadata.broker.list" or "bootstrap.servers" to be set with Kafka broker(s),
+  *   NOT zookeeper servers, specified in host1:port1,host2:port2 form.
+  * @param fromOffsets per-topic/partition Kafka offsets defining the (inclusive) starting point of the batch
+  * @param untilOffsets per-topic/partition Kafka offsets defining the (exclusive) ending point of the batch
+  * @param messageHandler function for translating each message into the desired type
+  */
+class KafkaRDD[
+  K: ClassTag,
+  V: ClassTag,
+  U <: Decoder[_]: ClassTag,
+  T <: Decoder[_]: ClassTag,
+  R: ClassTag](
+    sc: SparkContext,
+    kafkaParams: Map[String, String],
+    fromOffsets: Map[TopicAndPartition, Long],
+    untilOffsets: Map[TopicAndPartition, Long],
+    messageHandler: MessageAndMetadata[K, V] => R
+  ) extends RDD[R](sc, Nil) with Logging {
+
+  assert(fromOffsets.keys == untilOffsets.keys,
+    "Must provide both from and until offsets for each topic/partition")
+
+  override def getPartitions: Array[Partition] = fromOffsets.zipWithIndex.map { kvi =>
+    val ((tp, from), index) = kvi
+    new KafkaRDDPartition(index, tp.topic, tp.partition, from, untilOffsets(tp))
+  }.toArray
+
+  override def compute(thePart: Partition, context: TaskContext) = new NextIterator[R] {
+    context.addTaskCompletionListener{ context => closeIfNeeded() }
+
+    val part = thePart.asInstanceOf[KafkaRDDPartition]
+    val props = new Properties()
+    kafkaParams.foreach(param => props.put(param._1, param._2))
+    val fetchSize = Option(props.getProperty("fetch.message.max.bytes")).map(_.toInt).getOrElse(1024*1024)
+    val leaderBackoff = Option(props.getProperty("refresh.leader.backoff.ms")).map(_.toLong).getOrElse(200L)
+    val consumerConfig = new ConsumerConfig(props)
+    val keyDecoder = classTag[U].runtimeClass.getConstructor(classOf[VerifiableProperties])
+      .newInstance(consumerConfig.props)
+      .asInstanceOf[Decoder[K]]
+    val valueDecoder = classTag[T].runtimeClass.getConstructor(classOf[VerifiableProperties])
+      .newInstance(consumerConfig.props)
+      .asInstanceOf[Decoder[V]]
+    val consumer: SimpleConsumer = ???
+    var requestOffset = part.fromOffset
+    var iter: Iterator[MessageAndOffset] = null
+
+    override def getNext: R = {
+      if (iter == null || !iter.hasNext) {
+        val req = new FetchRequestBuilder().
+          addFetch(part.topic, part.partition, requestOffset, fetchSize).
+          build()
+        val resp = consumer.fetch(req)
+        if (resp.hasError) {
+          val err = resp.errorCode(part.topic, part.partition)
+          if (err == ErrorMapping.LeaderNotAvailableCode ||
+            err == ErrorMapping.NotLeaderForPartitionCode) {
+            log.error(s"Lost leader for topic ${part.topic} partition ${part.partition}, sleeping for ${leaderBackoff}ms")
+            Thread.sleep(leaderBackoff)
+          }
+          // Let normal rdd retry sort out reconnect attempts
+          throw ErrorMapping.exceptionFor(err)
+        }
+        iter = resp.messageSet(part.topic, part.partition)
+          .iterator
+          .filter(_.offset < requestOffset)
+      }
+      val item = iter.next
+      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))
+      }
+    }
+
+    override def close() = consumer.close()
+  }
+
+}