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

Posted 2023-03-31

【中文标题】GroupBy 和 Sum，如果大于零标记行 True

【英文标题】：GroupBy and Sum, If greater than zero mark row True

【发布时间】：2021-05-05 19:49:49

【问题描述】：

假设我有一个如下所示的数据集。我想按 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

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

【参考方案1】：

您可以对按 ID 分区的窗口进行求和，并检查总和是否为空（仅当该 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|
+---+-----+--------+