Pyspark：如何在数据框中复制一行 n 次？

Question

我有一个这样的数据框，如果列 n 大于一，我想复制行 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，但我不明白它是如何工作的......
谢谢

Answer 1

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"]) 

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

# 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))

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

# 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| 
+---+---+---+ 

Answer 2

我认为@Ahmed 的udf 回答是最好的方法，但这里有另一种方法，对于小型n 可能同样好或更好：

首先，在整个DataFrame上收集n的最大值：

max_n = df.select(f.max('n').alias('max_n')).first()['max_n']
print(max_n)
#3

现在为长度为max_n 的每一行创建一个数组，其中包含range(max_n) 中的数字。此中间步骤的输出将产生一个 DataFrame，如下所示：

df.withColumn('n_array', f.array([f.lit(i) for i in range(max_n)])).show()
#+---+---+---+---------+
#|  A|  B|  n|  n_array|
#+---+---+---+---------+
#|  1|  2|  1|[0, 1, 2]|
#|  2|  9|  1|[0, 1, 2]|
#|  3|  8|  2|[0, 1, 2]|
#|  4|  1|  1|[0, 1, 2]|
#|  5|  3|  3|[0, 1, 2]|
#+---+---+---+---------+

现在我们分解n_array 列，并过滤以仅保留数组中小于n 的值。这将确保我们拥有每行的n 副本。最后我们删除分解后的列得到最终结果：

df.withColumn('n_array', f.array([f.lit(i) for i in range(max_n)]))\
    .select('A', 'B', 'n', f.explode('n_array').alias('col'))\
    .where(f.col('col') < f.col('n'))\
    .drop('col')\
    .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|
#+---+---+---+

但是，我们正在为每一行创建一个 max_n 长度数组，而不是 udf 解决方案中的一个 n 长度数组。我目前还不清楚这将如何扩展与 udf 对于大型 max_n，但我怀疑 udf 会胜出。

Answer 3

使用 Spark 2.4.0+，使用内置函数更容易：array_repeat + explode：

from pyspark.sql.functions import expr

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

new_df = df.withColumn('n', expr('explode(array_repeat(n,int(n)))'))

>>> new_df.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|
+---+---+---+