pyspark sql 查询：根据条件计算不同的值

Question

我有一个如下的数据框：

+-----------+------------+-------------+-----------+
| id_doctor | id_patient | consumption | type_drug |
+-----------+------------+-------------+-----------+
| d1        | p1         |        12.0 | bhd       |
| d1        | p2         |        10.0 | lsd       |
| d1        | p1         |         6.0 | bhd       |
| d1        | p1         |        14.0 | carboxyl  |
| d2        | p1         |        12.0 | bhd       |
| d2        | p1         |        13.0 | bhd       |
| d2        | p2         |        12.0 | lsd       |
| d2        | p1         |         6.0 | bhd       |
| d2        | p2         |        12.0 | bhd       |
+-----------+------------+-------------+-----------+

我想计算每个医生服用 bhd 且消费 < 16.0 的不同患者。

我尝试了以下查询，但不起作用：

dataframe.groupBy(col("id_doctor"))
         .agg(
         countDistinct(col("id_patient")).where(col("type_drug") == "bhd" & col("consumption") < 16.0)
         )

有什么帮助吗？

谢谢！

Answer 1

PySpark 中的另一种解决方案，无需添加另一列：

dataframe.groupBy('id_doctor').agg(F.countDistinct(F.when(col("type_drug") == "bhd" & col("consumption") < 16.0, col('id_doctor')).otherwise(None)))