我认为@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()
然而,对于每一行我们正在创建一个长度为max_n
的数组- 而不是udf
解决方案中的长度为n
的数组。我不确定对于大型max_n
,这种方法如何才能够具有可伸缩性,但我认为udf
将获胜。