pyspark - 使用最大值为一列创建一个从 0 到该值的行值循环，并为其重复其他列值

Question

我有一个如下所示的数据框

Category   Descrption Max_value ... Total_months
 A         Keyboard   0              5
 B         Monitor    5              7
 

所以我期望的是低于从零到该值的总月数，其余列会重复。

 Category  Description Max_value ... months
  A        Keyborad    0               0
  A        Keyborad    0               1
  A        Keyboard    0               2
  A        Keyborad    0               3
  A        Keyborad    0               4
  A        Keyboard    0               5
  B        Monitor     5               0
  B        Monitor     5               1
  B        Monitor     5               2
  B        Monitor     5               3
  B        Monitor     5               4
  B        Monitor     5               5
  B        Monitor     5               6
  B        Monitor     5               7

为此，我在 SAS 中这样写

DATA FINAL_table;
SET INITIAL_table;
IF MAX_value NE . THEN DO;
DO months = 0 TO Total_months;

  OUTPUT;
END;
END;
ELSE DO;
months = .;
OUTPUT;
END;
RUN;   

如何在pyspark中实现这个dataframe？我想稍后对此进行绘制，因此每个月都需要一个数据点。请注意，这不是日历月，因此可能超过 12 个。我不想为此使用 pandas 或 numpy，因为以后转换变得很困难。

Answer 1

你可以生成一个序列，然后爆炸：

import pyspark.sql.functions as F

sequence = F.udf(lambda x: list(range(x+1)), 'array<int>')

df2 = df.withColumn('months', F.explode(sequence('Total_months')))