聚合具有类似"all"函数 pyspark 的 GroupBy 列

Question

我有一个包含主键、日期、变量和值的数据框架。我想按主键分组，并确定是否所有值都等于所提供的值。示例数据:

import pandas as pd
from datetime import date
from pyspark.sql import SparkSession
spark = SparkSession.builder.getOrCreate()
df = pd.DataFrame({
"pk": [1, 1, 1, 1, 2, 2, 2, 2, 3, 4],
"date": [
date("2022-05-06"),
date("2022-05-13"),
date("2022-05-06"),
date("2022-05-06"),
date("2022-05-14"),
date("2022-05-15"),
date("2022-05-05"),
date("2022-05-05"),
date("2022-05-11"),
date("2022-05-12")
],
"variable": [A, B, C, D, A, A, E, F, A, G],
"value": [2, 3, 2, 2, 1, 1, 1, 1, 5, 4]
})
df = spark.createDataFrame(df)
df.show()
df1.show()
#+-----+-----------+--------+-----+
#|pk   |       date|variable|value|
#+-----+-----------+--------+-----+
#|    1| 2022-05-06|       A|    2|
#|    1| 2022-05-13|       B|    3|
#|    1| 2022-05-06|       C|    2|
#|    1| 2022-05-06|       D|    2|
#|    2| 2022-05-14|       A|    1|
#|    2| 2022-05-15|       A|    1|
#|    2| 2022-05-05|       E|    1|
#|    2| 2022-05-05|       F|    1|
#|    3| 2022-05-11|       A|    5|
#|    4| 2022-05-12|       G|    4|
#+-----+-----------+--------+-----+

如果我想知道，给定一个主键pk，是否所有值都等于1(或任意布尔测试)，我该怎么做?我试过执行applyInPandas，但这不是超级高效，似乎可能有一个非常简单的方法来做到这一点。

Answer 1

+时，可以使用forall函数检查collect_list收集的所有值是否满足布尔检验。

import pyspark.sql.functions as F
df1 = (df
.groupby("pk")
.agg(F.expr("forall(collect_list(value), v -> v == 1)").alias("value"))
)
df1.show()
# +---+-----+
# | pk|value|
# +---+-----+
# |  1|false|
# |  3|false|
# |  2| true|
# |  4|false|
# +---+-----+
# or create a column using window function
df2 = df.withColumn("test", F.expr("forall(collect_list(value) over (partition by pk), v -> v == 1)"))
df2.show()
# +---+----------+--------+-----+-----+
# | pk|      date|variable|value| test|
# +---+----------+--------+-----+-----+
# |  1|2022-05-06|       A|    2|false|
# |  1|2022-05-13|       B|    3|false|
# |  1|2022-05-06|       C|    2|false|
# |  1|2022-05-06|       D|    2|false|
# |  3|2022-05-11|       A|    5|false|
# |  2|2022-05-14|       A|    1| true|
# |  2|2022-05-15|       A|    1| true|
# |  2|2022-05-05|       E|    1| true|
# |  2|2022-05-05|       F|    1| true|
# |  4|2022-05-12|       G|    4|false|
# +---+----------+--------+-----+-----+

你可能想把它放在case子句中来处理NULL值。