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)
    )
)

Answer 2

只需使用数据帧上的where - 此版本删除计数为 0 的id_doctor：

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

使用此语法，您可以保留所有"医生"：

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

Answer 3

和解决方案，无需添加额外的列（Scala）

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