[SPARK-23207][SQL] Shuffle+Repartition on a DataFrame could lead to incorrect answers

[jiangxb1987]

What changes were proposed in this pull request?

Currently shuffle repartition uses RoundRobinPartitioning, the generated result is nondeterministic since the sequence of input rows are not determined.

The bug can be triggered when there is a repartition call following a shuffle (which would lead to non-deterministic row ordering), as the pattern shows below:
upstream stage -> repartition stage -> result stage
(-> indicate a shuffle)
When one of the executors process goes down, some tasks on the repartition stage will be retried and generate inconsistent ordering, and some tasks of the result stage will be retried generating different data.

The following code returns 931532, instead of 1000000:

import scala.sys.process._

import org.apache.spark.TaskContext
val res = spark.range(0, 1000 * 1000, 1).repartition(200).map { x =>
  x
}.repartition(200).map { x =>
  if (TaskContext.get.attemptNumber == 0 && TaskContext.get.partitionId < 2) {
    throw new Exception("pkill -f java".!!)
  }
  x
}
res.distinct().count()

In this PR, we propose a most straight-forward way to fix this problem by performing a local sort before partitioning, after we make the input row ordering deterministic, the function from rows to partitions is fully deterministic too.

The downside of the approach is that with extra local sort inserted, the performance of repartition() will go down, so we add a new config named spark.sql.execution.sortBeforeRepartition to control whether this patch is applied. The patch is default enabled to be safe-by-default, but user may choose to manually turn it off to avoid performance regression.

This patch also changes the output rows ordering of repartition(), that leads to a bunch of test cases failure because they are comparing the results directly.

How was this patch tested?

Add unit test in ExchangeSuite.

With this patch(and spark.sql.execution.sortBeforeRepartition set to true), the following query returns 1000000:

import scala.sys.process._

import org.apache.spark.TaskContext

spark.conf.set("spark.sql.execution.sortBeforeRepartition", "true")

val res = spark.range(0, 1000 * 1000, 1).repartition(200).map { x =>
  x
}.repartition(200).map { x =>
  if (TaskContext.get.attemptNumber == 0 && TaskContext.get.partitionId < 2) {
    throw new Exception("pkill -f java".!!)
  }
  x
}
res.distinct().count()

res7: Long = 1000000

[SparkQA]

Test build #86635 has finished for PR 20393 at commit 7fd964e.

• This patch fails Spark unit tests.
• This patch merges cleanly.
• This patch adds the following public classes (experimental):
• public final class RecordBinaryComparator extends RecordComparator

[viirya]

retest this please.

[SparkQA]

Test build #86639 has finished for PR 20393 at commit 7fd964e.

• This patch passes all tests.
• This patch merges cleanly.
• This patch adds the following public classes (experimental):
• public final class RecordBinaryComparator extends RecordComparator

[gatorsmile]

LGTM

[shivaram]

@jiangxb1987 If I'm not wrong this problem will also happen with RDD repartition ? Will this fix also cover that ?

[sameeragarwal]

Yes, this bug also applies to RDD repartition but the current fix doesn't cover this (the local sort approach would be quite similar but it'll be a completely different codepath).

@jiangxb1987 - to @shivaram 's point, it'd be great to add a TODO for later.

[sameeragarwal]

LGTM, thanks!

[shivaram]

@sameeragarwal I think we should wait for the RDD fix for 2.3 as well ?

[jiangxb1987]

I added TODO on this, so we may have this for now and I'll continue working on the RDD path.

[SparkQA]

Test build #86668 has finished for PR 20393 at commit 400a766.

• This patch fails Spark unit tests.
• This patch merges cleanly.
• This patch adds no public classes.

[jiangxb1987]

retest this please

[SparkQA]

Test build #86682 has finished for PR 20393 at commit 400a766.

• This patch passes all tests.
• This patch merges cleanly.
• This patch adds no public classes.

[jiangxb1987]

Actually the similar approach cannot apply to fix RDD.repartition(), as in RDD[T], the data type T can be non-comparable, so we are not able to perform a local sort before actually repartition.

I'm stepping back investigating other approaches that requires some refactoring on the Core module but I don’t think that it is safe to ship the approach together with Spark 2.3

So my propose is, let’s include this PR in Spark 2.3, and target the follow up work to 2.4. Especially since the RDD.repartition() issue is not a regression of the latest version.

WDYT? @shivaram @sameeragarwal @rxin @mridulm

[jiangxb1987]

Another simple way to ensure correctness of RDD.repartition() is to do HashPartitioning instead of current RoundRobinPartitioning, but that will lead to regression when you have skew input data.

[mridulm]

@jiangxb1987 Other than hash partitioning, I dont see how this can be handled reliably ...
You are right, this is a basic correctness issue - unfortunately I never used this family of methods (coalasce, repartition, etc) and never saw the issue.

@shivaram Any thoughts ? You might have better insights.

[mridulm]

@jiangxb1987 Btw, we could argue this is a correctness issue since we added repartition - so not necessarily blocker :-)

[sameeragarwal]

Another (possibly cleaner) approach here would be to make the shuffle block fetch order deterministic but I agree that it might not be safe to include it in 2.3 this late.

edit: @jiangxb1987 how about calling this "... could lead to incorrect answers" instead of "data loss"?

[jiangxb1987]

Updated the title, does it sound good to have this PR? I'll open another one to address the RDD.repartition() issue (which will target to 2.4).

[sameeragarwal]

LGTM but we should get a broader consensus on this. In the meantime, I'm merging this patch to master/2.3.

[jiangxb1987]

I opened https://issues.apache.org/jira/browse/SPARK-23243 to track the RDD.repartition() patch, thanks for all the discussions! @shivaram @mridulm @sameeragarwal @gatorsmile

[shivaram]

I'm fine with merging this -- I just dont want to this issue to be forgotten for RDDs as I think its a major correctness issue.

@mridulm @sameeragarwal Lets continue the discussion on the new JIRA.

[mridulm]

@sameeragarwal I am not sure if we can make shuffle fetch deterministic - without quite a lot of perf overhead; do you have any thoughts on how to do this in case I am missing something here ?

[sameeragarwal]

@mridulm one approach that Xingbo is looking into (independently of #20414) is to have the ShuffleBlockFetcherIterator remember the order of blocks it fetches and store them in that order. Given that the blocks will still be fetched in parallel, depending on the available buffer size, we'll then have to spill some out-of-order blocks on disk in order to avoid OOMs on the receiver (similar to #16989). While this would still regress performance, it might be better than the current local sort based fix. Note that I'm not arguing against the fact that hash partitioning would be the "best" fix in terms of performance, but it'd then defeat the purpose of repartition (due to skew).

[tcondie]

@jiangxb1987 and @sameeragarwal we are seeing this issue in Spark 2.3.2 when a cache step is introduced after each repartition operation. I have not been able to repro it using the example listed in this PR and Jira. Could either of you please verify that this bug fix is complete and that adding a cache step would not affect the solution?

[tcondie]

@jiangxb1987 could you please provide a little guidance on how to run the example repro for this issue? Spark seems to fail the job entirely when the kill switch brings down the executor, and consequently triggers the need to re-execute upstream stages, which Spark seems to punt on i.e., not do.

[sc-abhisheksoni]

I encountered this bug in my code on databricks:
20-50 Node cluster was using - Spark 2.4.3, Scala 2.11

I had a Dataframe with records for Unique IDs. Once the data was ready to write to BLOB storage, I repartitioned the data to 2 partitions and wrote to storage.

On reading the records back from storage the number of total records remained the same but the number of Unique Ids reduced. On looking at the data the repartitioning introduced duplicate records in the data that was written to BLOB.
As descriibed above the problem was non deterministic, and sometimes we got correct number of unique records while other times some data was duplicated.

Once I removed Repartitioning I have not encountered this issue again.

[jiangxb1987]

@jiangxb1987 could you please provide a little guidance on how to run the example repro for this issue? Spark seems to fail the job entirely when the kill switch brings down the executor, and consequently triggers the need to re-execute upstream stages, which Spark seems to punt on i.e., not do.

IIRC When I run this example on databricks notebook, I started a cluster with 20 workers, and set spark.stage.maxConsecutiveAttempts to a very big number so the failed stage keeps retry.