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

我有一个数据框如下：

+-----------+------------+-------------+-----------+
| 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 的消耗量

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

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

有什么帮助吗？

thanks!

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

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