如何在pyspark中重命名列，类似于使用与Spark兼容的SQL PIVOT语句?

Question

>我在Spark中运行了这个查询

f"""
select
col_cate,
true as segment_true,
false as segment_false
from (
SELECT
{feature},
{target_col},
count(*) as cnt
from
table1
group by
{feature},
{target_col}
) pivot (
sum(cnt)
for target_bool in (true, false)
)
""")

输入数据是这个

+--------+-----------+
|col_cate|target_bool|
+--------+-----------+
|       A|       true|
|       A|      false|
|       B|       true|
|       B|      false|
|       A|       true|
|       A|      false|
|       B|       true|
|       B|      false|
+--------+-----------+

输出数据为

+--------+------------+-------------+
|col_cate|segment_true|segment_false|
+--------+------------+-------------+
|       A|        true|        false|
|       B|        true|        false|
+--------+------------+-------------+

但是，当我尝试在pyspark中执行相同的操作时，我无法弄清楚如何将输出从[col_cate, true, false]重命名为[col_cate, segment_true segment_false]

我该怎么做？

我试过了

output_df.groupBy(["col_cate", "target_bool"]).
count().
groupBy("col_cate").
pivot("target_bool").
sum("count")

但是似乎没有办法重命名代码中的列。我知道我可以在之后重命名它，但它感觉不那么优雅。

Answer 1

您可以相应地转换列 如

df.withColumn('target_bool', when(col('target_bool')=='true',lit('segment_true')).otherwise(lit('segment_false'))).
groupBy(["col_cate", "target_bool"]).
count().
groupBy("col_cate").
pivot("target_bool").
sum("count").show()

+--------+-------------+------------+
|col_cate|segment_false|segment_true|
+--------+-------------+------------+
|       B|            2|           2|
|       A|            2|           2|
+--------+-------------+------------+

或者您的 SQL 等效版本为

output_df.groupBy(["col_cate", "target_bool"]).
count().
groupBy("col_cate").
pivot("target_bool").
sum("count").select(col('col_cate'),col('true').alias('segment_true'),col('false').alias('segment_false'))

Answer 2

您可以在透视后使用 agg 和 alias 方法执行此操作。

import pyspark.sql.functions as F
tst = sqlContext.createDataFrame([('A','true',1),('A','false',2),('B','true',3),('B','false',4),('A','true',5),('A','false',6),('B','true',7),('B','false',8)],schema=['col_cate','target_bool','id'])
#%%
tst_res =tst.groupBy(["col_cate", "target_bool"]).count().groupBy("col_cate").pivot("target_bool").agg(F.sum("count").alias("segment"),F.sum("count").alias("dummy"))

结果：

+--------+-------------+-----------+------------+----------+
|col_cate|false_segment|false_dummy|true_segment|true_dummy|
+--------+-------------+-----------+------------+----------+
|       B|            2|          2|           2|         2|
|       A|            2|          2|           2|         2|
+--------+-------------+-----------+------------+----------+

第二次聚合的原因，假人我也不清楚。但是，仅当存在多个聚合时，才会进行重命名。我正在调查这个。但这应该对你有用。

编辑 ：如评论中所述，您无需将其分组两次。下面的代码将达到目的：tst_res1=tst.groupby("col_cate").pivot("target_bool").agg(F.count("target_bool").alias('segment'),F.count("target_bool").alias('dummy'))