How does Spark Handles Partitions and Shuffles

Diving deep into Spark (PySpark) for the first time, and I'm already flooded with under-the-hood curiosity.

From the docs, I understand that when creating an RDD, Spark splits the data evenly across the available partitions. For example...

rdd = spark.sparkContext.parallelize([("A", 1), ("A", 2), ("A", 3), ("B", 1), ("B", 2), ("C", 1), ("C", 2), ("C", 3)])

So my understanding is that if there is 4 partitions then the following is true...

P1: ("A", 1), ("A", 2)
P2: ("A", 3), ("B", 1)
P3: ("B", 2), ("C", 1)
P4: ("C", 2), ("C", 3)

QUESTION 1: I get that each partition will always get 2 elements, but will they always get the SAME two elements? i.e If I run the rdd code above 1 million times, will the elements belong to the same partition every time?

My next question has to do with groupByKey(). Say we are finding the sum for each key..

sums = rdd.groupByKey().mapValues(sum)

Here is where Spark loses me. I know there are still 4 partitions but not sure how the data is distributed after the shuffle caused by the grouping.

QUESTION 2: How does Spark decide where to send records? And what do the partitions look like after the shuffle?

Here are a few different possible scenarios I can think of after the grouping and sum transformations...

SCENARIO 1 (EVENLY)
P1: ("A", 6)
P2: ("B", 3)
P3: ("C", 6)
P4: EMPTY

SCENARIO 2 (SKEWED)
P1: ("A", 6), ("B", 3)
P2: ("C", 6)
P3: EMPTY
P4: EMPTY

SCENARIO 3 (REALLY SKEWED AND OUT OF ORDER)
P1: EMPTY
P2: EMPTY
P3: ("A", 6), ("B", 3), ("C", 6)
P4: EMPTY

I'm hoping Spark has a process it always follows for distributing data after shuffles. If anyone has any answers I'd love to hear them.

Thanks!

The number of partitions can be specified by a parameter of parallelize(). Otherwise, it is defaultParallelism. For example, assume that you are using local mode, then local[4] will result in 4 partitions.

Question 1

> I get that each partition will always get 2 elements, but will they always get the SAME two elements?

When the number of partitions is fixed, how elements are assigned within partitions are determined by different partitioning strategies according to the types of elements. By default, it is partitioned roughly evenly which is implemented simply based on sequences' indices. Although the code is written in Scala, I think you should have no trouble of understanding it:

def positions(length: Long, numSlices: Int): Iterator[(Int, Int)] = {
  (0 until numSlices).iterator.map { i =>
    val start = ((i * length) / numSlices).toInt
    val end = (((i + 1) * length) / numSlices).toInt
    (start, end)
  }
}

Therefore, given 8 elements with 4 partitions, each partition has 2 elements. No matter how many times you run the code, how elements are partitioned is determined.

> I know there are still 4 partitions but not sure how the data is distributed after the shuffle caused by the grouping.

The signature of groupByKey() is:

def groupByKey(self: RDD[tuple[K, V]],
               numPartitions: int | None = None,
               partitionFunc: (K) -> int = portable_hash) -> RDD[tuple[K, Iterable[V]]]

It hash-partitions the resulting RDD with numPartitions (4 in this case) partitions. Given an element (k, v), it will be shuffled to partition hash(k) % 4.

Question 2

I think the answer above also fits Question 2. Note that developers can also specify user-defined Partitioner based on Range/Hash methods.

Given an RDD with 4 partitions like this:

<!-- language-all: scala -->

val rdd = sc.parallelize(Seq(("A", 1), ("A", 2), ("A", 3), ("B", 1), ("B", 2), ("C", 1), ("C", 2), ("C", 3)), 4)

You might always get the same result irrespective of how many instances you run of the same line because you're using the parallelize method with a small sample of data. How data gets partitioned when it is read depends on a lot of factors like the underlying file systems, the type of files being read, number of executors, number of driver cores etc.

So to answer your question 1, no the partition structure will not necessarily remain the same everytime you read the data.

For question 2, whenever you used key based operations, a shuffle is introduced which moves partitions with the same key in the same partition.

For example the following code:

val sumRdd = rdd.groupByKey.mapValues(_.sum)

Would still return an RDD with 4 partitions. The data will be moved based on the default partitioning scheme - Hash Partitioning which decides the partition key based on the logic:

object.hashCode % numPartitions

So all objects with the same hashCode will move to the same partition and you will have an underlying structure like:

sumRdd.mapPartitionsWithIndex{ (idx, itr) => itr.toList.map( c => c+" -> partition#"+idx ).iterator }.collect
// Array[String] = Array((A,6) -> partition#1, (B,3) -> partition#2, (C,6) -> partition#3)

sumRdd.getNumPartitions will still return 4 but since one partition is empty, no tasks will be invoked for that partition.