Как использовать DataFrame.withColumn с условием

Question

Я хочу знать, как создать новый столбец в полном DataFrame, но значения которого основаны только на подмножестве DataFrame (т. Е. Некоторые функции применяются на основе условия).

В этом примере я хочу создать столбец, который связывает каждый идентификатор с суммой только положительных значений в столбце Значение по идентификатору.

Использование этого набора данных:

+-----+-----+
|  ID |Value|
+-----+-----+
|  1  |  -4 |
+-----+-----+
|  2  |  5  |
+-----+-----+
|  2  | -23 |
+-----+-----+
|  1  |  5  |
+-----+-----+
|  2  | 19  |
+-----+-----+
|  1  | 12  |
+-----+-----+

Я хочу получить следующий результат:

+-----+-----+-------------+
|  ID |Value| sum_pos_val |
+-----+-----+-------------+
|  1  |  -4 |     17      |
+-----+-----+-------------+
|  2  |  5  |     24      |
+-----+-----+-------------+
|  2  | -23 |     24      |
+-----+-----+-------------+
|  1  |  5  |     17      |
+-----+-----+-------------+
|  2  | 19  |     24      |
+-----+-----+-------------+
|  1  | 12  |     17      |
+-----+-----+-------------+

Следующий код будет суммировать столбец Значение по идентификатору, но как я могу отфильтровать только положительные значения для суммирования по?

import pyspark.sql.functions as f
from pyspark.sql import Window

w = Window.partitionBy('ID')
df.withColumn('sum_pos_val', f.sum('Value').over(w))\
    .show()

Answer 1

Заменить

f.sum('Value').over(w)

на

f.sum(
    f.when(f.col('Value') > 0, f.col('Value')).otherwise(0)
).over(w)

или еще проще (см. Группировка искровых данных в кадре не учитывает нули )

f.sum(
    f.when(f.col('Value') > 0, f.col('Value'))
).over(w)