您有两个解决方案。
第一个是可以在
col3
上使用数据透视表获取唯一值的计数,然后将您的数据透视表与聚合的数据框连接起来,计算其他列的总和/平均数/最小值/最大值。
代码如下:
from pyspark.sql import functions as F
result = df \
.groupBy('col1', 'col2') \
.pivot('col3') \
.agg(F.count('col3')) \
.fillna(0) \
.join(
df.groupby('col1', 'col2').agg(F.sum('col4').alias('col4_sum')),
['col1', 'col2']
)
使用您提供的数据框,您将获得以下结果:
+----+----+---+---+---+--------+
|col1|col2|a |b |c |col4_sum|
+----+----+---+---+---+--------+
|1 |1 |1 |1 |0 |3 |
|1 |2 |0 |0 |1 |1 |
|2 |1 |1 |1 |0 |7 |
+----+----+---+---+---+--------+
然而,你无法选择从数据透视表中提取的列的名称,它将是值的名称。
如果你真的想要选择列的名称,可以先检索所有不同的值,然后从每个值构建你的汇总列,具体如下:
from pyspark.sql import functions as F
values = map(lambda x: x.col3, df.select("col3").distinct().collect())
count_of_distinct_values = [F.sum((F.col('col3') == i).cast('integer')).alias('count_' + i) for i in values]
other_column_aggregations = [F.sum('col4').alias('col4_sum')]
aggregated = count_of_distinct_values + other_column_aggregations
result = df.groupBy('col1', 'col2').agg(*aggregated)
你将获得以下数据框:
+
|col1|col2|count_a|count_b|count_c|col4_sum|
+
|1 |1 |1 |1 |0 |3 |
|1 |2 |0 |0 |1 |1 |
|2 |1 |1 |1 |0 |7 |
+