Python PySpark DataFrame randomSplit方法用法及代码示例

例子

考虑以下PySpark DataFrame：

df = spark.createDataFrame([["Alex", 20], ["Bob", 30], ["Cathy", 40], ["Dave", 40]], ["name", "age"])
df.show()



+-----+---+
| name|age|
+-----+---+
| Alex| 20|
|  Bob| 30|
|Cathy| 40|
| Dave| 40|
+-----+---+

将 PySpark 数据帧随机拆分为更小的 DataFrames

将此 PySpark DataFrame 随机拆分为 2 sub-DataFrames，行拆分为 75-25：

list_dfs = df.randomSplit([0.75,0.25], seed=24)
for _df in list_dfs:
    _df.show()



+-----+---+
| name|age|
+-----+---+
| Alex| 20|
|Cathy| 40|
+-----+---+

+----+---+
|name|age|
+----+---+
| Bob| 30|
|Dave| 40|
+----+---+

尽管我们预计第一个 DataFrame 包含 3 行，而第二个 DataFrame 包含 1 行，但我们看到分割是 50-50。这是因为，如上所述，randomSplit(~) 基于伯努利采样。

种子参数的怪癖

seed 参数用于再现性。例如，考虑以下PySpark DataFrame：

df = spark.createDataFrame([["Alex", 20], ["Bob", 30], ["Cathy", 40], ["Dave", 40]], ["name", "age"])
df



+-----+---+
| name|age|
+-----+---+
| Alex| 20|
|  Bob| 30|
|Cathy| 40|
| Dave| 40|
+-----+---+

鉴于 PySpark DataFrame 以完全相同的方式分区，使用相同的种子运行 randomSplit(~) 方法将保证相同的分割：

list_dfs = df.randomSplit([0.75,0.25], seed=24)
for _df in list_dfs:
    _df.show()



+-----+---+
| name|age|
+-----+---+
| Alex| 20|
|Cathy| 40|
+-----+---+

+----+---+
|name|age|
+----+---+
| Bob| 30|
|Dave| 40|
+----+---+

多次运行上述操作将始终产生相同的分割，因为 PySpark DataFrame 的分区是相同的。

我们可以通过将 DataFrame 转换为 RDD，然后使用 glom() 方法来查看 PySpark DataFrame 的行是如何分区的：

df = spark.createDataFrame([["Alex", 20], ["Bob", 30], ["Cathy", 40], ["Dave", 40]], ["name", "age"])
df.rdd.glom().collect()



[[],
 [Row(name='Alex', age=20)],
 [],
 [Row(name='Bob', age=30)],
 [],
 [Row(name='Cathy', age=40)],
 [],
 [Row(name='Dave', age=40)]]

在这里，我们看到 PySpark DataFrame 被分为 8 个分区，但其中一半是空的。

让我们使用 repartition(~) 更改分区：

df = df.repartition(2)
df.rdd.glom().collect()



[[Row(name='Alex', age=20),
  Row(name='Bob', age=30),
  Row(name='Cathy', age=40),
  Row(name='Dave', age=40)],
 []]

尽管 DataFrame 的内容相同，但我们现在只有 2 个分区，而不是 8 个分区。

让我们使用与之前相同的种子 ( 24 ) 来调用 randomSplit(~) ：

list_dfs = df.randomSplit([0.75,0.25], seed=24)
for _df in list_dfs:
    _df.show()



+-----+---+
| name|age|
+-----+---+
| Alex| 20|
|  Bob| 30|
|Cathy| 40|
| Dave| 40|
+-----+---+

+----+---+
|name|age|
+----+---+
+----+---+

请注意，即使我们使用相同的种子，我们最终也会得到不同的分裂。这证实了 seed 参数仅在底层分区相同时才保证一致的分割。您应该谨慎对待这种行为，因为分区在洗牌操作后可能会发生变化(例如 join(~) 和 groupBy(~) )。