[SPARK-23207][SPARK-22905][SPARK-24564][SPARK-25114][SQL][BACKPORT-2.2] Shuffle+Repartition on a DataFrame could lead to incorrect answers

[bersprockets]

What changes were proposed in this pull request?

Back port of #20393.

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.

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

Author: Xingbo Jiang [email protected]

How was this patch tested?

Ran all SBT unit tests for org.apache.spark.sql.*.

Ran pyspark tests for module pyspark-sql.

[gatorsmile]

cc @jiangxb1987

[gatorsmile]

The original fix efccc02 has been merged to Spark 2.3. After 5+ months, we have not received any correctness regression that is caused by this fix, although this fix definitely will introduce a performance regression. I think we should merge it. The general risk is not very high and it resolves a serious correctness bug.

[SparkQA]

Test build #94619 has finished for PR 22079 at commit efccc02.

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

[bersprockets]

The test "model load / save" in ChiSqSelectorSuite fails because of this line in
ChiSqSelector.scala

spark.createDataFrame(dataArray).repartition(1).write.parquet(Loader.dataPath(path))

In 2.4, the line is:

spark.createDataFrame(sc.makeRDD(dataArray, 1)).write.parquet(Loader.dataPath(path))

If you change 2.4 to also have that line, and also remove the follow-up PR (#20426) to avoid sorting when there is one partition, this test also fails on 2.4 in the same way.

So I am not sure which way to go: Update ChiSqSelector.scala to be like 2.4 (simply a one line change), or make the test accept this new order.

[SparkQA]

Test build #94633 has finished for PR 22079 at commit 495cba5.

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

[jiangxb1987]

We shall also include #20088 in this backport PR.

[bersprockets]

@jiangxb1987

We shall also include #20088 in this backport PR.

I did that shortly after commenting, which allowed the tests to pass. I squashed it into the first commit, so it wasn't obvious I did it.

Should I also include #20426 in this PR, or treat that separately?

[jiangxb1987]

Both seems fine to me, it's just a minor improvement. Normally we don't backport a improvement, but since it's a simple and small change I'm confident it is safe to also include the change in a backport PR.

[SparkQA]

Test build #94656 has finished for PR 22079 at commit 8d2d558.

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

[bersprockets]

Hmmm... I somehow managed to break SparkR tests but fixing a comment. It seems to have auto-retried and broke the second time too.

[bersprockets]

retest this please

[SparkQA]

Test build #94665 has finished for PR 22079 at commit 8d2d558.

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

[mridulm]

It is possible for two objects to be unequal and yet we consider them as equal with this code, if the long values are separated by Int.MaxValue.
Any particular reason why this code is written like this with the inaccuracy and not simply the idiomatic comparator comparison ?

[squito]

good question, @jiangxb1987 any thoughts on this? (it was this way in the original)

[jiangxb1987]

As far as I can recall It's actually a mistake. At first the function returned a Long value and later I changed the return value to Integer, let me fix it. Thanks for discovering this @mridulm !

[dongjoon-hyun]

@jiangxb1987 and @bersprockets . SPARK-22905 consists of two commits.

• ChiSqSelector ([SPARK-22905][MLLIB] Fix ChiSqSelectorModel save implementation #20088)
• GaussianMixtureModel ([SPARK-22905][ML][FollowUp] Fix GaussianMixtureModel save #20113)

If we want to include SPARK-22905 here, it had better be explicit and complete by putting [SPARK-22905] into the PR title and includes both patches.

[bersprockets]

@dongjoon-hyun Thanks, I will include the other commit and also update the title.

[SparkQA]

Test build #94701 has finished for PR 22079 at commit 81f57fe.

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

[SparkQA]

Test build #94704 has finished for PR 22079 at commit bab8e68.

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

[squito]

overall I'm in favor of backporting this, and it looks like the only changes to the original were very small, so I'm in favor of this.

[bersprockets]

@jiangxb1987

Here are some of the differences from the original PR

• I also ported the follow up PR [SPARK-23207][SQL][FOLLOW-UP] Don't perform local sort for DataFrame.repartition(1) #20426
• I ported [SPARK-22905][MLLIB] Fix ChiSqSelectorModel save implementation #20088 (for SPARK-22905) to get the tests to pass. I also ported its followup, [SPARK-22905][ML][FollowUp] Fix GaussianMixtureModel save #20113.
• sql/catalyst/src/main/java/org/apache/spark/sql/execution/UnsafeExternalRowSorter.java line 112
  UnsafeExternalSorter.create does not take a Supplier for the 5th argument, so I put get() on the argument to directly pass a RecordComparator object.
• sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchange.scala: I access SQLConf differently (since SQLConf.get doesn't appear to work in 2.2).

[bersprockets]

Once this is merged, I will also back-port:

• [SPARK-24564][TEST] Add test suite for RecordBinaryComparator
• [SPARK-25114][Core] Fix RecordBinaryComparator when subtraction between two words is divisible by Integer.MAX_VALUE. #22101

[bersprockets]

@jiangxb1987 gentle ping.

[gatorsmile]

The PR #22101 has been merged. Please merge the latest one to this PR. Thanks!

[bersprockets]

@gatorsmile So I should include all the related PRs merged to master as a single PR here? Just verifying.

[gatorsmile]

ok to test

[SparkQA]

Test build #95053 has finished for PR 22079 at commit 2edad85.

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

[SparkQA]

Test build #95054 has finished for PR 22079 at commit 2edad85.

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

[gatorsmile]

cc @jiangxb1987

[jiangxb1987]

LGTM, thanks!

[gatorsmile]

Thanks! Merged to 2.2

[bersprockets]

@gatorsmile Weird, I don't see it on branch-2.2. Is that a sync issue?

[gatorsmile]

d7c3aae Done

[gatorsmile]

@bersprockets Could you please close this PR?