pyspark - 将字符串列的 Spark 数据框拆分为多个 bool 列

Question

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

   id | value 
------+--------
    0 | A,B
    1 | A,C
    2 | B

我们想把它改造成:

   id | A     | B     | C
------+-------+-------+-------+
    0 | True  | True  | False | 
    1 | True  | False | True  | 
    2 | False | True  | False | 

进行这种转换的最佳方法是什么？

Answer 1

假设这是您的输入数据框:

df = spark.createDataFrame([(0,["A","B"]),(1,["A","C"]),(2, ["B"])],["id","value"])

然后使用 explode 和 pivot 得到一个包含整数和空值的表。

df2 = df.withColumn("x",explode(df.value)).drop("value").groupBy("id").pivot("x").count()
df2.show()

+---+----+----+----+
| id|   A|   B|   C|
+---+----+----+----+
|  0|   1|   1|null|
|  1|   1|null|   1|
|  2|null|   1|null|
+---+----+----+----+

最后，你只需要将值转换为 bool 值，例如:

for col_name in df2.columns[1:]:
    df2 = df2.withColumn(col_name, col(col_name).isNotNull())

df2.show()

+---+-----+-----+-----+
| id|    A|    B|    C|
+---+-----+-----+-----+
|  0| true| true|false|
|  1| true|false| true|
|  2|false| true|false|
+---+-----+-----+-----+