Sending Large CSV to Kafka using python Spark

Question

I am trying to send a large CSV to kafka. The basic structure is to read a line of the CSV and zip it with the header.

a = dict(zip(header, line.split(",")

This then gets converted to a json with:

message = json.dumps(a)

I then use kafka-python library to send the message

from kafka import SimpleProducer, KafkaClient
kafka = KafkaClient("localhost:9092")
producer = SimpleProducer(kafka)
producer.send_messages("topic", message)

Using PYSPARK I have easily created an RDD of messages from the CSV file

sc = SparkContext()
text = sc.textFile("file.csv")
header = text.first().split(',')
def remove_header(itr_index, itr):
    return iter(list(itr)[1:]) if itr_index == 0 else itr
noHeader = text.mapPartitionsWithIndex(remove_header)

messageRDD = noHeader.map(lambda x: json.dumps(dict(zip(header, x.split(","))

Now I want to send these messages: I define a function

def sendkafka(message):
  kafka = KafkaClient("localhost:9092")
  producer = SimpleProducer(kafka)
  return producer.send_messages('topic',message)

Then I create a new RDD to send the messages

sentRDD = messageRDD.map(lambda x: kafkasend(x))

I then call sentRDD.count()

Which starts churning and sending messages

Unfortunately this is very slow. It sends 1000 messages a second. This is on a 10 node cluster of 4 cpus each and 8gb of memory.

In comparison, creating the messages takes about 7 seconds on a 10 million row csv. ~ about 2gb

I think the issue is that I am instantiating a kafka producer inside the function. However, if I don't then spark complains that the producer doesn't exist even though I have tried defining it globally.

Perhaps someone can shed some light on how this problem may be approached.

Thank you,

Answer 1

You can create a single producer per partition and use either mapPartitions or foreachPartition:

def sendkafka(messages):
    kafka = KafkaClient("localhost:9092")
    producer = SimpleProducer(kafka)
    for message in messages:
        yield producer.send_messages('topic', message)

sentRDD = messageRDD.mapPartitions(sendkafka)

If above alone won't help you can try to extend it using an asynchronous producer.

In Spark 2.x it is also possible to use Kafka data source. You'll have to include spark-sql-kafka jar, matching Spark and Scala version (here 2.2.0 and 2.11 respectively):

spark.jars.packages  org.apache.spark:spark-sql-kafka-0-10_2.11:2.2.0

convert data to a DataFrame (if it is not DataFrame already):

messageDF = spark.createDataFrame(messageRDD, "string")

and write using DataFrameWriter:

(messageDF.write
    .format("kafka")
    .option("topic", topic_name)
    .option("kafka.bootstrap.servers", bootstrap_servers)
    .save())