按列进行Spark重新分区，每列具有动态分区数

Question

如何根据一列中项目的数量对DataFrame进行分区。假设我们有一个可容纳100个人的DataFrame（列为first_name和country），并且我们想为一个国家/地区的每10个人创建一个分区。

如果我们的数据集包含来自中国的80个人，来自法国的15个人和来自古巴的5个人，那么我们将需要8个用于中国的分区，2个用于法国的分区和1个用于古巴的分区。

这是无效的代码：

• df.repartition($"country")：这将为中国创建一个分区，为法国创建一个分区，为古巴创建一个分区
• df.repartition(8, $"country", rand)：这将为每个国家/地区最多创建8个分区，因此应为中国创建8个分区，但法国和古巴的分区未知。法国可能在8个分区中，而古巴则可能在5个分区中。有关更多详细信息，请参见this answer。

这是repartition()文档：

当我查看repartition()方法时，我什至看不到带有三个参数的方法，因此似乎其中某些行为未得到记录。

有没有办法动态设置每一列的分区数？它将使创建分区数据集的方式更加容易。

Answer 1

由于spark分区数据的方式，您将无法完全实现这一目标。 Spark使用您在重新分区中指定的列，将该值哈希为64b长，然后根据分区数对值求模。这样，分区的数量是确定的。之所以如此工作，是因为联接除了要确保两端的哈希相同之外，还需要联接左右两侧的分区数量匹配。

“我们想为一个国家的每10个人创建一个分区。”

您到底想在这里完成什么？一个分区中只有10行可能会降低性能。您是否要创建一个分区表，其中分区中的每个文件都被保证只有x个行？

“ df.repartition（$“ country”）：这将为中国创建1个分区，为法国创建1个分区，为古巴创建1个分区”

这实际上将创建一个数据框，其中包含按国家/地区哈希的默认洗牌分区数量

  def repartition(partitionExprs: Column*): Dataset[T] = {
    repartition(sparkSession.sessionState.conf.numShufflePartitions, partitionExprs: _*)
  }

“ df.repartition（8，$” country“，rand）：每个国家最多可以创建8个分区，因此应该为中国创建8个分区，但是法国和古巴的分区是未知的。在8个分区中，古巴最多可以在5个分区中。有关更多详细信息，请参阅此答案。”

就像明智的做法，这是完全错误的。只有8个分区，而这8个分区中的国家基本上是随机洗牌的。

Answer 2

下面的代码将为每个数据文件（sample dataset is here）创建十行：

val outputPath = new java.io.File("./tmp/partitioned_lake5/").getCanonicalPath
df
  .repartition(col("person_country"))
  .write
  .option("maxRecordsPerFile", 10)
  .partitionBy("person_country")
  .csv(outputPath)

这是Spark 2.2之前的代码，每个数据文件将创建大约10行：

val desiredRowsPerPartition = 10

val joinedDF = df
  .join(countDF, Seq("person_country"))
  .withColumn(
    "my_secret_partition_key",
    (rand(10) * col("count") / desiredRowsPerPartition).cast(IntegerType)
  )

val outputPath = new java.io.File("./tmp/partitioned_lake6/").getCanonicalPath
joinedDF
  .repartition(col("person_country"), col("my_secret_partition_key"))
  .drop("count", "my_secret_partition_key")
  .write
  .partitionBy("person_country")
  .csv(outputPath)

This blog post详细介绍了partitioning。