apache-spark - 从一列有条件地计数

Question

我有一些数据存储在 df1 :

| id1 | id2      | product |
|-----|----------|---------|
| 1   | abc-2323 | Upload  |
| 1   | 234234   | Upload  |
| 1   | 43322    | Upload  |
| 2   | abc-449  | Upload  |
| 3   | abc-495  | Upload  |
| 3   | 432      | Upload  |
| 3   | 6543     | Upload  |

每 id1和 product ，我想计算不同的 id2 s 基于天气，它们以 abc- 开头或不。我尝试过这样的条件聚合:

agg_data = (
    df1
    .groupby('id1', 'product')
    .agg(
       sf.when(~ sf.col('id2').like('abc-%'), sf.countDistinct('id2')).alias('id2_count_without_abc'),
       sf.when(sf.col('id2').like('abc-%'), sf.countDistinct('id2')).alias('is2_count_with_abc')
    )
)

然而，这会出错 id2既不存在于 group by 中，也不是聚合函数。我不确定为什么我不能做这个有条件的，因为只是 sf.countDistinct('id2')会自己工作。例如。这没有问题:

agg_data = (
    df1
    .groupby('id1', 'product')
    .agg(
       sf.countDistinct('id2').alias('id2_count_without_abc'),
       sf.countDistinct('id2').alias('is2_count_with_abc')
    )
)

Answer 1

当您这样做时groupBy每组你会得到很多行。因为您可以为 countDistinct() 获得每组一个值，你不会得到错误。但是对于 when 子句，您可以一次传递一个值，但 group 将返回多行。
要消除此错误，您可以使用 first() 选择一个值或 last() (这将返回相应组的第一个/最后一个值)并将其传递给 when()功能。
但是您需要 group 的所有值才能根据一个条件进行计数计算。下面的代码应该可以解决问题。

df1.groupby('id1', 'product') \
.agg(sf.collect_list('id2').alias('id2')) \
.withColumn('id2_count_with_abc', sf.size(sf.expr("filter(id2, i-> i like 'abc-%')"))) \
.withColumn('id2_count_without_abc', sf.abs(sf.size(sf.col('id2')) - sf.col('id2_count_with_abc'))) \
.drop('id2').show()


+---+-------+------------------+---------------------+
|id1|product|id2_count_with_abc|id2_count_without_abc|
+---+-------+------------------+---------------------+
|  2| Upload|                 1|                    0|
|  1| Upload|                 1|                    2|
|  3| Upload|                 1|                    2|
+---+-------+------------------+---------------------+