https://github.com/apache/spark
Revision 18e83f9ab234e2e475842c70f191d90e4dd6e00d authored by Bobby Wang on 22 September 2022, 12:59:00 UTC, committed by Wenchen Fan on 22 September 2022, 13:00:11 UTC
### What changes were proposed in this pull request?

``` scala
val df = spark.range(0, 100, 1, 50).repartition(4)
val v = df.rdd.mapPartitions { iter => {
        Iterator.single(iter.length)
}.collect()
println(v.mkString(","))
```

The above simple code outputs `50,0,0,50`, which means there is no data in partition 1 and partition 2.

The RoundRobin seems to ensure to distribute the records evenly *in the same partition*, and not guarantee it between partitions.

Below is the code to generate the key

``` scala
      case RoundRobinPartitioning(numPartitions) =>
        // Distributes elements evenly across output partitions, starting from a random partition.
        var position = new Random(TaskContext.get().partitionId()).nextInt(numPartitions)
        (row: InternalRow) =>
{         // The HashPartitioner will handle the `mod` by the number of partitions
         position += 1
         position
 }
```

In this case, There are 50 partitions, each partition will only compute 2 elements. The issue for RoundRobin here is it always starts with position=2 to do the Roundrobin.

See the output of Random
``` scala
scala> (1 to 200).foreach(partitionId => print(new Random(partitionId).nextInt(4) + " "))  // the position is always 2.
2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2
```

Similarly, the below Random code also outputs the same value,

``` scala
(1 to 200).foreach(partitionId => print(new Random(partitionId).nextInt(2) + " "))
(1 to 200).foreach(partitionId => print(new Random(partitionId).nextInt(4) + " "))
(1 to 200).foreach(partitionId => print(new Random(partitionId).nextInt(8) + " "))
(1 to 200).foreach(partitionId => print(new Random(partitionId).nextInt(16) + " "))
(1 to 200).foreach(partitionId => print(new Random(partitionId).nextInt(32) + " "))
```

Consider partition 0, the total elements are [0, 1], so when shuffle writes, for element 0, the key will be (position + 1) = 2 + 1 = 3%4=3, the element 1, the key will be (position + 1)=(3+1)=4%4 = 0
consider partition 1, the total elements are [2, 3], so when shuffle writes, for element 2, the key will be (position + 1) = 2 + 1 = 3%4=3, the element 3, the key will be (position + 1)=(3+1)=4%4 = 0

The calculation is also applied for other left partitions since the starting position is always 2 for this case.

So, as you can see, each partition will write its elements to Partition [0, 3], which results in Partition [1, 2] without any data.

This PR changes the starting position of RoundRobin. The default position calculated by `new Random(partitionId).nextInt(numPartitions)` may always be the same for different partitions, which means each partition will output the data into the same keys when shuffle writes, and some keys may not have any data in some special cases.

### Why are the changes needed?

The PR can fix the data skew issue for the special cases.

### Does this PR introduce _any_ user-facing change?

No

### How was this patch tested?

Will add some tests and watch CI pass

Closes #37855 from wbo4958/roundrobin-data-skew.

Authored-by: Bobby Wang <wbo4958@gmail.com>
Signed-off-by: Wenchen Fan <wenchen@databricks.com>
(cherry picked from commit f6c4e58b85d7486c70cd6d58aae208f037e657fa)
Signed-off-by: Wenchen Fan <wenchen@databricks.com>
1 parent 3ac1c3d
History
Tip revision: 18e83f9ab234e2e475842c70f191d90e4dd6e00d authored by Bobby Wang on 22 September 2022, 12:59:00 UTC
[SPARK-40407][SQL] Fix the potential data skew caused by df.repartition
Tip revision: 18e83f9

README.md

back to top