根据 pyspark 中的 groupby 过滤行创建具有最大值的新列

Question

I have a spark dataframe

import pandas as pd
foo = pd.DataFrame({'id': [1,1,2,2,2], 'col': ['a','b','a','a','b'], 'value': [1,5,2,3,4],
'col_b': ['a','c','a','a','c']})

I want to create a new column with the max of the value column, groupped by id . But I want the max value only for the rows that col==col_b

My result spark dataframe should look like this

foo = pd.DataFrame({'id': [1,1,2,2,2], 'col': ['a','b','a','a','b'], 'value': [1,5,2,3,4],
'max_value':[1,1,3,3,3], 'col_b': ['a','c','a','a','c']})

I have tried

from pyspark.sql import functions as f
from pyspark.sql.window import Window
w = Window.partitionBy('id')
foo = foo.withColumn('max_value', f.max('value').over(w))\
    .where(f.col('col') == f.col('col_b'))

But I end up losing some rows.

Any ideas?

Answer 1

Use when function for conditionnal aggregation max :

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

w = Window.partitionBy('id')

foo = foo.withColumn('max_value', F.max(F.when(F.col('col') == F.col('col_b'), F.col('value'))).over(w))