0

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

  1. 这两件事都发生在他们身上，
  2. 有 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 中优雅地实现这一目标？

4

1 回答 1

0

使用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           |
+-----+------------+------------+
于 2021-06-22T07:54:40.883 回答