GroupBy 和 Sum，如果大于零标记行 True

Question

假设我有一个如下所示的数据集。我想按 ID 分组，如果所有值都为 NULL，或者当总和不大于 0 时创建一个新列，将每一行标记为 True 或 False。

下面是理想的数据集。

ID  VALUE
1    6.00
1    5.00
2    NULL
2    NULL
2    NULL
3    3.00

ID  VALUE   ALL_NULL
1    6.00    False
1    5.00    False
2    NULL    True
2    NULL    True
2    NULL    True
3    3.00    False

这是我目前为止尝试过的方法。

df = (
    wbcl_full
    .withColumn('ALL_NULL', F.when((wbcl_full.groupBy('ID').agg(sum('VALUE').collect()[0][0]) == 0), True)
    .otherwise(False))
)
return df

感谢您的帮助！理想情况下，我还没有找到执行条件 window 函数的好方法。在某些情况下，当我 window 根据分组结果结束时，我想有选择地选择要选择的行。

Answer 1

您可以对按 ID 分区的 window 求和，并检查总和是否为空（仅当该 ID 的所有元素都为空时才会发生）或者它是否不大于 0：

from pyspark.sql import functions as F, Window

result = df.withColumn(
    'ALL_NULL', 
    F.sum('VALUE').over(Window.partitionBy('ID')).isNull() | 
    (F.sum('VALUE').over(Window.partitionBy('ID')) <= 0)
)

result.show()
+---+-----+--------+
| ID|VALUE|ALL_NULL|
+---+-----+--------+
|  1|  6.0|   false|
|  1|  5.0|   false|
|  3|  3.0|   false|
|  2| null|    true|
|  2| null|    true|
|  2| null|    true|
+---+-----+--------+

GroupBy 和 Sum，如果大于零标记行 True

GroupBy and Sum, If greater than zero mark row True

window-functions

apache-spark

apache-spark-sql

pyspark