Question

我想知道如何在完整的DataFrame上创建新列，但其值仅基于DataFrame的子集（即，根据条件应用的某些函数）。

在此示例中，我想创建一个列，该列将每个ID与ID中“值”列中仅正值之和相关联。

使用的是数据集：

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

以下代码将按ID对“值”列进行求和，但是如何过滤出仅对正值求和呢？

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)

或更简单的（见spark dataframe groupping does not count nulls）

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

如何在条件下使用DataFrame.withColumn

1 个答案: