从 PySpark 中的虚拟列中获取单列

Question

我有一个如下所示的 PySpark df，其中事件 2020-01 和 2020-02 是假人（我总共有 18 个月）。

df = (
    sc.parallelize([
        ("A", 1, 0, 1), ("B", 0, 1, 0), ("C", 0, 1, 0),
        ("D", 1, 1, 1),
    ]).toDF(["id", "event", "2020-01", "2020-02"])
)

id  event   2020-01   2020-02
A   1       0         1
B   0       1         0
C   0       1         0
D   1       1         1

我想创建一个带有 id、事件和月份的新 df，其中月份是从 "2020-01" 和 "2020-02" 创建的列，其中这些列 == 1。所以所需的 df 如下所示：

id  event   month
A   1       2020-02
B   0       2020-01
C   0       2020-01
D   1       2020-01
D   1       2020-02

我找到了一个熊猫解决方案here，它可以满足我的需求，但我的 df 对于熊猫来说太大了。我没有设法让this 解决方案起作用。它只是为每个 ID 每个月创建。

Answer 1

如果您有很多列，这可能会起作用：

out_list = [i for i in df.columns if i not in ['id','event']]

stack_str = ','.join(map(lambda x:'"{0}",`{0}`'.format(x),out_list))

(df.selectExpr('id','event',
               'stack({0},{1}) as 
               (feature,value)'.format(len(out_list),stack_str))
 .where(col('value')==1)
 .show()[enter image description here][1]
)

Answer 2

您可以通过以下方式做到这一点

from spark.sql import functions

df1 = df.select("id", "event").where(df["2020-01"] == 1).withColumn("month", functions.lit("2020-01"))
df2 = df.select("id", "event").where(df["2020-02"] == 1).withColumn("month", functions.lit("2020-02"))

df1 = df1.unionAll(df2).orderBy("id")
df1.show()
+---+-----+-------+
| id|event|  month|
+---+-----+-------+
|  A|    1|2020-02|
|  B|    0|2020-01|
|  C|    0|2020-01|
|  D|    1|2020-01|
|  D|    1|2020-02|
+---+-----+-------+