有条件地控制求和中的值

Question

我有一个数据框架，看起来像这样:

| id   | c1  | c2   | c3    |
|------|-----|------|-------|
| 1334 | 20  | 3565 | 0.005 |
| 1335 | 543 | 2100 | 0.205 |

c3由c1 / (c1 + c2)计算，如下所示:

agg = (
df1
.groupby('id')
.agg(
F.count('c1').alias('c1'),
F.count('c2').alias('c2')
)
).withColumn('c3',
F.col('c1') / (F.col('c1') + F.col('c2')))

我想有条件地改变我的.withColumn的值，像这样:

if c1 < 50 then 0
if c2 > 1000 then 1000

对于id1334，计算结果为0 / (0 + 1000)，对于id1335，计算结果为545 / (543 + 1000)。

我曾尝试使用.when()，但似乎无法获得语法正确

Answer 1

你可以试试:

.withColumn(
'c3',
F.when(F.col('c1') < 50, 0).otherwise(F.col('c1')) / (
F.when(F.col('c1') < 50, 0).otherwise(F.col('c1')) +
F.when(F.col('c2') > 1000, 1000).otherwise(F.col('c2'))
)
)