基于在pyspark中使用groupby过滤的行创建具有最大值的新列

Question

我有一个spark 数据框

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']})

我想用value 列的max 创建一个新列，按id 分组。但我只希望 max value 用于 col==col_b 的行

我的结果 spark 数据框应该是这样的

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']})

我试过了

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'))

但我最终会丢失一些行。

有什么想法吗？

Answer 1

使用when函数进行条件聚合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))