python - pyspark:如何在数据帧中重复n次？

我有一个这样的数据框，如果列n大于1，我想将行重复n次:

A   B   n
1   2   1
2   9   1
3   8   2
4   1   1
5   3   3

并像这样转换:

A   B   n
1   2   1
2   9   1
3   8   2
3   8   2
4   1   1
5   3   3
5   3   3
5   3   3

我想我应该使用explode，但是我不明白它是如何工作的...
谢谢

最佳答案

explode函数为给定数组或映射中的每个元素返回一个新行。

利用此功能的一种方法是使用udf为每行创建一个大小为n的列表。然后爆炸生成的数组。

from pyspark.sql.functions import udf, explode
from pyspark.sql.types import ArrayType, IntegerType

df = spark.createDataFrame([(1,2,1), (2,9,1), (3,8,2), (4,1,1), (5,3,3)] ,["A", "B", "n"])

# use udf function to transform the n value to n times
n_to_array = udf(lambda n : [n] * n, ArrayType(IntegerType()))
df2 = df.withColumn('n', n_to_array(df.n))

# now use explode
df2.withColumn('n', explode(df2.n)).show()

+---+---+---+
| A | B | n |
+---+---+---+
|  1|  2|  1|
|  2|  9|  1|
|  3|  8|  2|
|  3|  8|  2|
|  4|  1|  1|
|  5|  3|  3|
|  5|  3|  3|
|  5|  3|  3|
+---+---+---+

关于python - pyspark:如何在数据帧中重复n次？，我们在Stack Overflow上找到一个类似的问题：https://stackoverflow.com/questions/50624745/