pyspark sql query: подсчитать разные значения с условиями

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 для каждого врача.</p>

Я пробовал следующий запрос, но он не работает:

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

любая помощь?

спасибо!

Answer 1

Просто используйте where на вашем фрейме данных - эта версия удаляет id_doctor, где счетчик равен 0:

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