在 pyspark 中，如何在 groupby 之后在一列中找到值，而不是在另一列（设置减法）中查找值？答案

【问题标题】：In pyspark, how do I find values in one column but not in another (set subtraction) after groupby?在 pyspark 中，如何在 groupby 之后在一列中找到值，而不是在另一列（设置减法）中查找值？
【发布时间】：2021-09-05 12:31:37
【问题描述】：

我有一个 pyspark datframe df 的记录，每条记录都有id 和group，并标记是否发生了两个事件（event1, event2）。我想找到每个组中的 id 数量，即：

这两件事都发生在他们身上，
他们有 event2 但没有 event1。

我在这里提取一个简单的例子：

df:
|  id | event1 | event2 | group
| 001 |      1 |      0 |     A
| 001 |      1 |      0 |     A    
| 001 |      1 |      1 |     A  
| 002 |      0 |      1 |     A  
| 003 |      1 |      0 |     A  
| 003 |      1 |      1 |     A  
| ... |    ... |    ... |     B
...

在上面的df中，对于group = A，有2个id有event1:(001,003)，3个id有event2:(001,002,003)。因此，例如，event2 而不是 event1 中的 id 数为 1。

我希望得到这样的东西。

group | event2_not_1 | event1_and_2 |
    A |            1 |            2 |
    B |          ... |          ... |

到目前为止，我已经尝试收集为每个事件出现的一组 id，然后在new_df 中分别执行集合操作。但我觉得这很笨拙。例如，

df_new = (
  df.withColumn('event1_id', when(col('event1') == 1, col('id')))
    .withColumn('event2_id', when(col('event2') == 1, col('id')))
    .groupby('group').agg(collect_set('event1_id').alias('has_event1'),
                          collect_set('event2_id').alias('has_event2'))
)

如何在 pyspark 中优雅地实现这一点？

【问题讨论】：

标签： pyspark set-operations

【解决方案1】：

使用groupby 两次。

df.groupBy("group", "id").agg(f.max("event1").alias("event1"), f.max("event2").alias("event2")) \
  .groupBy("group").agg(f.sum(f.expr("if(event2 = 1 and event1 = 0, 1, 0)")).alias("event2_not_1"), f.sum(f.expr("if(event1 = 1 and event2 = 1, 1, 0)")).alias("event1_and_2"))

+-----+------------+------------+
|group|event2_not_1|event1_and_2|
+-----+------------+------------+
|A    |1           |2           |
+-----+------------+------------+

【讨论】：