Optimizing Spark Direct Kafka Consumption: Subpartition Concurrency and Repartition Strategies

When using Apache Spark's Direct approach to read from Kafka, some Spark partitions finish early while others lag, extending batch execution time or causing missed consumption deadlines.

The community identified this issue as SPARK‑22056 and proposed modifying the KafkaRDD class's getPartitions method to introduce a configurable topic.partition.subconcurrency parameter, allowing multiple Spark partitions to consume a single Kafka partition.

Original implementation:

override def getPartitions: Array[Partition] = {
  offsetRanges.zipWithIndex.map {
    case (o, i) => 
      val (host, port) = leaders(TopicAndPartition(o.topic, o.partition))
      new KafkaRDDPartition(i, o.topic, o.partition, o.fromOffset, o.untilOffset, host, port)
  }.toArray
}

Modified implementation (adds sub‑concurrency logic):

override def getPartitions: Array[Partition] = {
  val subconcurrency = if (kafkaParams.contains("topic.partition.subconcurrency"))
    kafkaParams.getOrElse("topic.partition.subconcurrency","1").toInt
    else 1
  val numPartitions = offsetRanges.length
  val subOffsetRanges = new Array[OffsetRange](subconcurrency * numPartitions)
  for (i <- 0 until numPartitions) {
    val offsetRange = offsetRanges(i)
    val step = (offsetRange.untilOffset - offsetRange.fromOffset) / subconcurrency
    var from = offsetRange.fromOffset - 1L
    var until = offsetRange.fromOffset - 1L
    for (j <- 0 until subconcurrency) {
      from = offsetRange.fromOffset + j * step
      until = if (j == subconcurrency - 1) offsetRange.untilOffset
              else offsetRange.fromOffset + (j + 1) * step - 1
      subOffsetRanges(i * subconcurrency + j) = OffsetRange.create(offsetRange.topic, offsetRange.partition, from, until)
    }
  }
  subOffsetRanges.zipWithIndex.map {
    case (o, i) =>
      val (host, port) = leaders(TopicAndPartition(o.topic, o.partition))
      new KafkaRDDPartition(i, o.topic, o.partition, o.fromOffset, o.untilOffset, host, port)
  }.toArray
}

Setting topic.partition.subconcurrency to 1 preserves the original one‑to‑one mapping; values greater than 1 split a Kafka partition across multiple Spark partitions, improving throughput but potentially breaking ordering guarantees.

Community feedback highlighted concerns about data ordering and violating the design principle that each Kafka partition maps to a single Spark partition.

Alternative approaches include explicitly repartitioning or coalescing the RDD after creation, e.g.:

val kafkaStream = KafkaUtils.createDirectStream[String, String, StringDecoder, StringDecoder](
  streamingContext, kafkaParams, topics).repartition(xxx).mapPartitions(xxx)

This preserves order within each logical key group but adds overhead; it is beneficial only when the repartitioning cost is outweighed by the gain in parallelism.

Another suggestion is to spawn multiple threads inside mapPartitions to process records concurrently, though this does not address skew caused by uneven Kafka partition sizes.

Ultimately, handling Kafka‑induced data skew may require adjusting Kafka partition counts, improving partitioning strategy, or scaling Spark resources, rather than relying on multi‑threaded consumption within a single partition.