通过减少其中一列来复制行答案

【问题标题】：Replicate rows by decrementing one of the columns通过减少其中一列来复制行
【发布时间】：2021-07-06 03:20:18
【问题描述】：

我有一个包含几列的 Pyspark 数据框。为简单起见，我们只考虑两列 id 和 num：

id	num
1001	5
1002	3

这里的目标是为id 列创建多行，但减少num 列。以下是预期的输出：

id	num
1001	1
1001	2
1001	3
1001	4
1001	5
1002	1
1002	2
1002	3

当前解决方案：

创建一个 UDF，为所有缺失值创建一个列表，如下所示：

id	num
1001	[1,2,3,4,5]
1002	[1,2,3]

然后使用explode 创建那些复制的行。

有没有更好的方法来实现这一点？除非绝对必要，否则我不想使用 UDF。

【问题讨论】：

标签： python apache-spark pyspark apache-spark-sql

【解决方案1】：

您可以使用sequence 生成从文字 1 到 num 值的数字范围

import pyspark.sql.functions as F

df.withColumn('num', F.explode(F.sequence(F.lit(1), F.col('num'))))

【讨论】：

那是完美的。正是我想要的。谢谢！

【解决方案2】：

给定：

+----+---+
|  id|num|
+----+---+
|1001|  5|
|1002|  3|
+----+---+

如果您想重复行num 次，您可以使用Pandas_udf。您不一定需要创建带有数组的 num 列来复制行。

from pyspark.sql.functions import monotonically_increasing_id
def repeat_rows(iterator:Iterator[pd.DataFrame])->Iterator[pd.DataFrame]:
  for pdf in iterator:
    yield(pdf.loc[pdf.index.repeat(pdf['num'])].reset_index(drop=True))
df.mapInPandas(repeat_rows, df.schema).withColumn('num', row_number().over(Window.partitionBy(col('id')).orderBy(col('id')))).show()

+----+---+
|  id|num|
+----+---+
|1001|  1|
|1001|  2|
|1001|  3|
|1001|  4|
|1001|  5|
|1002|  1|
|1002|  2|
|1002|  3|
+----+---+

【讨论】：

糟糕，请参阅我的编辑。我错过了一些东西。希望它为您添加新方法