Условный контроль значений в сумме

#python #apache-spark #pyspark #apache-spark-sql

#python #apache-spark #pyspark #apache-spark-sql

Вопрос:

У меня есть фрейм данных, который выглядит следующим образом:

 | 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
 

Таким образом, для id 1334 вычисление становится 0 / (0 1000) , а для id 1335 оно становится 545 / (543 1000) .

Я пытался использовать .when() , но, похоже, не могу получить правильный синтаксис

Ответ №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'))
    )
)