RDDs: The API Underneath DataFrames

RDDs — Resilient Distributed Datasets — were Spark’s original abstraction. DataFrames are built on them, and understanding RDDs explains several DataFrame behaviours that otherwise look arbitrary.

Creating one

// From a local collection
val nums = sc.parallelize(1 to 20, 4)
println(s"partitions: ${nums.getNumPartitions}")
println(s"first five: ${nums.take(5).mkString(", ")}")

// From a file — one element per line
val lines = sc.textFile("data/shakespeare.txt")
println(s"lines: ${lines.count()}")
println(s"first: ${lines.first()}")

partitions: 4
first five: 1, 2, 3, 4, 5
lines: 124456
first: THE SONNETS

An RDD is a partitioned collection with no schema. Spark knows it holds Int or String objects and nothing more — which is precisely why it cannot optimise operations on it.

Transformations and actions

The same lazy split as DataFrames:

val words = lines
  .flatMap(_.toLowerCase.split("\\W+"))
  .filter(_.nonEmpty)

println("nothing has run yet")

val counts = words
  .map(w => (w, 1))
  .reduceByKey(_ + _)

counts.sortBy(-_._2).take(5).foreach { case (w, n) => println(f"$w%-8s $n%6d") }

nothing has run yet
the       27843
and       26847
i         20538
to        19822
of        18192

flatMap, filter, map and reduceByKey are transformations. take is the action that triggered everything.

reduceByKey versus groupByKey

The single most important RDD performance rule:

import java.lang.System.nanoTime

def time[T](label: String)(f: => T): T = {
  val t0 = nanoTime(); val r = f
  println(f"$label%-24s ${(nanoTime() - t0) / 1e9}%6.2fs"); r
}

val pairs = words.map(w => (w, 1))

time("reduceByKey") { pairs.reduceByKey(_ + _).count() }
time("groupByKey")  { pairs.groupByKey().mapValues(_.sum).count() }

reduceByKey                0.94s
groupByKey                 4.71s

Five times slower for an identical result. reduceByKey sums within each partition first, so one partial count per word per partition crosses the network. groupByKey shuffles every single occurrence of every word, then groups.

On a hot key groupByKey does not just run slowly — it fails:

java.lang.OutOfMemoryError: Java heap space
	at org.apache.spark.util.collection.CompactBuffer.growToSize(CompactBuffer.scala:143)

All values for one key must fit in one executor’s memory. reduceByKey never materialises that list.

RDD versus DataFrame, measured

Same word count, both APIs:

import org.apache.spark.sql.functions._

time("RDD word count") {
  lines.flatMap(_.toLowerCase.split("\\W+")).filter(_.nonEmpty)
       .map(w => (w, 1)).reduceByKey(_ + _).count()
}

val df = spark.read.text("data/shakespeare.txt")
time("DataFrame word count") {
  df.select(explode(split(lower(col("value")), "\\W+")).as("word"))
    .filter(length(col("word")) > 0)
    .groupBy("word").count().count()
}

RDD word count             0.94s
DataFrame word count       0.38s

Two and a half times faster, from three things RDDs cannot do: Catalyst reordered and pruned the plan, Tungsten stored rows in a compact off-heap format, and whole-stage code generation compiled the operators into one Java method.

Inspect the compiled code:

df.select(explode(split(lower(col("value")), "\\W+")).as("word"))
  .groupBy("word").count()
  .queryExecution.debug.codegen()

Found 2 WholeStageCodegen subtrees.
== Subtree 1 / 2 (maxMethodCodeSize:2847; maxConstantPoolSize:312(0.48% used); numInnerClasses:0) ==
*(1) HashAggregate(keys=[word#12], functions=[partial_count(1)])
...
/* 042 */   private void agg_doAggregateWithKeys_0() throws java.io.IOException {
/* 043 */     while ( inputadapter_input_0.hasNext()) {

Spark generated Java source for your query and compiled it. An RDD closure is an opaque function object — there is nothing to generate.

Partitioning control

This is where RDDs still earn their place. You can supply a custom partitioner:

import org.apache.spark.Partitioner

class RegionPartitioner(regions: Array[String]) extends Partitioner {
  private val index = regions.zipWithIndex.toMap
  override def numPartitions: Int = regions.length
  override def getPartition(key: Any): Int = index(key.asInstanceOf[String])
}

val events = sc.parallelize(Seq(
  ("uk", 12), ("us", 40), ("de", 7), ("uk", 3), ("us", 18), ("de", 22)
))

val partitioned = events.partitionBy(new RegionPartitioner(Array("uk", "us", "de")))

partitioned.mapPartitionsWithIndex { (i, it) =>
  Iterator(s"partition $i: ${it.toList.mkString(", ")}")
}.collect().foreach(println)

partition 0: (uk,12), (uk,3)
partition 1: (us,40), (us,18)
partition 2: (de,7), (de,22)

Exact placement, chosen by you. The DataFrame API offers repartition(col) but always uses hash partitioning — it cannot give you “UK goes to partition 0”.

Partitioning is also sticky, which avoids repeated shuffles:

val cached = partitioned.cache()
println(s"partitioner: ${cached.partitioner}")

time("first reduceByKey")  { cached.reduceByKey(_ + _).count() }
time("second reduceByKey") { cached.reduceByKey(_ + _).count() }

partitioner: Some(RegionPartitioner@4a2fb1c8)
first reduceByKey          0.31s
second reduceByKey         0.02s

Because the RDD already knows its keys are correctly placed, reduceByKey needs no shuffle at all — it aggregates within each partition.

Per-partition work

mapPartitions runs your function once per partition rather than once per element, which matters when there is expensive setup:

val ids = sc.parallelize(1 to 100000, 8)

time("map — connect per element") {
  ids.map { id => Thread.sleep(0); s"row-$id" }.count()
}

time("mapPartitions — per partition") {
  ids.mapPartitions { it =>
    // Imagine: open a database connection here, once per partition
    it.map(id => s"row-$id")
  }.count()
}

map — connect per element        1.84s
mapPartitions — per partition    0.21s

Opening a connection inside map opens it 100,000 times. Inside mapPartitions it opens 8 times. This pattern has no clean DataFrame equivalent outside pandas UDFs.

Converting between the two

val df2 = spark.createDataFrame(Seq(("a", 1), ("b", 2))).toDF("k", "v")

val asRdd = df2.rdd
println(s"as RDD: ${asRdd.first()}   type: ${asRdd.first().getClass.getSimpleName}")

import org.apache.spark.sql.types._
import org.apache.spark.sql.Row
val schema = StructType(Seq(StructField("k", StringType), StructField("v", IntegerType)))
val backToDf = spark.createDataFrame(asRdd, schema)
backToDf.show()

as RDD: [a,1]   type: GenericRowWithSchema
+---+---+
|  k|  v|
+---+---+
|  a|  1|
|  b|  2|
+---+---+

Each conversion is a full pass and discards Catalyst’s knowledge of the plan. Round-tripping in the middle of a pipeline is a common accidental performance cliff.

When to reach for an RDD

The default is DataFrame. Drop to RDDs deliberately, for one of the reasons above, and convert back as soon as the data has structure.

Practice

reduceByKey  shuffle write: 1.2 MB
groupByKey   shuffle write: 47.8 MB

first:  0.31s
second: 0.02s

map:           83.2s   (setup once per element)
mapPartitions:  0.6s   (setup once per partition)

before: *(1) Filter (amount#3 > 500.0)   [whole-stage codegen]
after:  Scan ExistingRDD[...]            [no codegen, no pushdown]

Next: the SQL interface, and where Catalyst does its work.